#include "cuda_helper.h"
#include "hdr_synthesis.h"

#include <cuda_profiler_api.h>

#include <nppi_color_conversion.h>
#include <nppi_filtering_functions.h>

#include <opencv2/core/cuda.hpp>
#include <opencv2/imgcodecs.hpp>

#include <boost/iostreams/device/mapped_file.hpp>

#include <iostream>
#include <chrono>

static constexpr auto image_width = 2448;
static constexpr auto image_height = 2048;

cv::Mat download_image(void *ptr, size_t pitch, size_t width, size_t height, int type) {
    auto gpu_mat = cv::cuda::GpuMat{(int) height, (int) width, type, ptr, pitch};
    cv::Mat mat;
    gpu_mat.download(mat);
    return mat;
}

struct image_buffer {
    void *pyr_image_f32, *pyr_weight_f32;
    size_t image_pitch, weight_pitch;
};

int main() {
    auto path_a = "/home/tpx/project/HDRSynthesis/data/phantom2_5ms.raw";
    auto path_b = "/home/tpx/project/HDRSynthesis/data/phantom2_20ms.raw";

    using boost::iostreams::mapped_file;
    auto img_file_a = mapped_file{path_a, boost::iostreams::mapped_file_base::readonly};
    auto img_file_b = mapped_file{path_b, boost::iostreams::mapped_file_base::readonly};

    auto hdr = hdr_synthesizer{image_width, image_height};
    void *buf_a, *buf_b;
    hdr.malloc_buffer(&buf_a);
    hdr.malloc_buffer(&buf_b);
    hdr.preprocess_image(buf_a, (uint8_t *) img_file_a.const_data());
    hdr.preprocess_image(buf_b, (uint8_t *) img_file_b.const_data());

    auto img_hdr_dev = cv::cuda::GpuMat{image_height, image_width, CV_8UC3};
    hdr.merge_image(buf_b, buf_a, img_hdr_dev.cudaPtr(), img_hdr_dev.step1());

//    for (int i = 0; i < 8; ++i) {
//        auto start_ts = std::chrono::system_clock::now();
//        cudaDeviceSynchronize();
//        hdr.preprocess_image(buf_a, (uint8_t *) img_file_a.const_data());
//        hdr.preprocess_image(buf_b, (uint8_t *) img_file_b.const_data());
//        cudaDeviceSynchronize();
//        std::cout << std::chrono::duration_cast<std::chrono::microseconds>(
//                std::chrono::system_clock::now() - start_ts).count() << std::endl;
//    }
//
//    cudaProfilerStart();
//    hdr.preprocess_image(buf_a, (uint8_t *) img_file_a.const_data());
//    cudaProfilerStop();
//

    auto real_ptr = (image_buffer *) buf_b;
    auto host_rgb_a = download_image(img_hdr_dev.cudaPtr(), img_hdr_dev.step1(),
                                     image_width, image_height, CV_8UC3);
//    auto host_rgb_a = download_image((char *) real_ptr->pyr_image_f32, real_ptr->image_pitch,
//                                     image_width, image_height, CV_32FC3);
//    void *ptr;
//    size_t pitch;
//    hdr.test_func(&ptr, &pitch);
//    auto host_rgb_a = download_image(ptr, pitch,
//                                     image_width, image_height, CV_32FC3);

    double min_val, max_val;
    cv::minMaxLoc(host_rgb_a, &min_val, &max_val);
    std::cout << min_val << " " << max_val << " " << cv::mean(host_rgb_a) << std::endl;

    cv::imwrite("test.bmp", host_rgb_a);

    return 0;
}