#include <jni.h>
#include <string>
#include <vector>
// #include "ncnn/yolov5ncnn.h"

#include <opencv2/opencv.hpp>
#include <opencv2/core.hpp>
#include <opencv2/imgproc.hpp>


inline std::string jstring2string(JNIEnv *env, jstring jStr)
{
    if (!jStr)
        return "";
    const jclass stringClass = env->GetObjectClass(jStr);
    const jmethodID getBytes = env->GetMethodID(stringClass, "getBytes", "(Ljava/lang/String;)[B");
    const jbyteArray stringJbytes = (jbyteArray) env->CallObjectMethod(jStr, getBytes, env->NewStringUTF("UTF-8"));
    size_t length = (size_t) env->GetArrayLength(stringJbytes);
    jbyte* pBytes = env->GetByteArrayElements(stringJbytes, NULL);
    std::string ret = std::string((char *)pBytes, length);
    env->ReleaseByteArrayElements(stringJbytes, pBytes, JNI_ABORT);
    env->DeleteLocalRef(stringJbytes);
    env->DeleteLocalRef(stringClass);
    return ret;
}

bool makeHdr(std::vector<float>& times, std::vector<std::string>& paths, cv::Mat& rgb)
{
    // Read images and exposure times
    std::vector<cv::Mat> images;

    for (auto it = paths.cbegin(); it != paths.cend(); ++it)
    {
        cv::Mat im = cv::imread((*it).c_str());
        images.push_back(im);
    }
    // Align input images
    // cout << "Aligning images ... " << endl;
    cv::Ptr<cv::AlignMTB> alignMTB = cv::createAlignMTB();
#if 0
    alignMTB->process(images, images);
#endif

    // Obtain Camera Response Function (CRF)
    // cout << "Calculating Camera Response Function (CRF) ... " << endl;
    cv::Mat responseDebevec;
    cv::Ptr<cv::CalibrateDebevec> calibrateDebevec = cv::createCalibrateDebevec();
    calibrateDebevec->process(images, responseDebevec, times);

    // Merge images into an HDR linear image
    // cout << "Merging images into one HDR image ... ";
    cv::Mat hdrDebevec;
    cv::Ptr<cv::MergeDebevec> mergeDebevec = cv::createMergeDebevec();
    mergeDebevec->process(images, hdrDebevec, times, responseDebevec);
    // Save HDR image.
    // imwrite((OUTPUT_DIR "hdrDebevec.hdr"), hdrDebevec);
    // cout << "saved hdrDebevec.hdr " << endl;

    {
        std::vector<cv::Mat> empty;
        empty.swap(images);
    }

    // Tonemap using Reinhard's method to obtain 24-bit color image
    // cout << "Tonemaping using Reinhard's method ... ";
    cv::Mat ldrReinhard;
    cv::Ptr<cv::TonemapReinhard> tonemapReinhard = cv::createTonemapReinhard(1.5, 0, 0, 0);
    tonemapReinhard->process(hdrDebevec, ldrReinhard);
    hdrDebevec.release();

    int type = ldrReinhard.type();
    ldrReinhard = ldrReinhard * 255;

    ldrReinhard.convertTo(rgb, CV_8U);
    ldrReinhard.release();

    return true;
}

extern "C"
JNIEXPORT jboolean JNICALL
Java_com_xypower_camera2raw_Camera2RawFragment_makeHdr__JLjava_lang_String_2JLjava_lang_String_2Ljava_lang_String_2(
        JNIEnv *env, jobject thiz, jlong exposureTime1, jstring path1, jlong exposureTime2,
        jstring path2, jstring outputPath) {

    std::vector<float> times;
    std::vector<std::string> paths;

    times.push_back((float)(exposureTime1 / 1000) / 1000000.0);
    times.push_back((float)(exposureTime2 / 1000) / 1000000.0);

    paths.push_back(jstring2string(env, path1));
    paths.push_back(jstring2string(env, path2));

    cv::Mat rgb;
    if (makeHdr(times, paths, rgb))
    {
        std::string fileName = jstring2string(env, outputPath);
        if (cv::imwrite(fileName.c_str(), rgb))
        {
            return JNI_TRUE;
        }
    }

    return JNI_FALSE;
}

extern "C"
JNIEXPORT jboolean JNICALL
Java_com_xypower_camera2raw_Camera2RawFragment_makeHdr__JLjava_lang_String_2JLjava_lang_String_2JLjava_lang_String_2Ljava_lang_String_2(
        JNIEnv *env, jobject thiz, jlong exposureTime1, jstring path1, jlong exposureTime2,
        jstring path2, jlong exposureTime3, jstring path3, jstring outputPath) {

    std::vector<float> times;
    std::vector<std::string> paths;

    times.push_back((float)(exposureTime1 / 1000) / 1000000.0);
    times.push_back((float)(exposureTime2 / 1000) / 1000000.0);
    times.push_back((float)(exposureTime3 / 1000) / 1000000.0);

    paths.push_back(jstring2string(env, path1));
    paths.push_back(jstring2string(env, path2));
    paths.push_back(jstring2string(env, path3));

    cv::Mat rgb;
    if (makeHdr(times, paths, rgb))
    {
        std::string fileName = jstring2string(env, outputPath);
        if (cv::imwrite(fileName.c_str(), rgb))
        {
            return JNI_TRUE;
        }
    }

    return JNI_FALSE;
}