#include <fstream>
#include <iostream>
#include <stdio.h>

#ifdef _DEBUG
#define STB_IMAGE_WRITE_IMPLEMENTATION
#include <include/stb_image_write.h>
#endif

#include <hdrplus_pipeline.h>
#include "Burst.h"

#include <include/HDRPlus.h>

extern "C" void halide_register_argv_and_metadata(
        int (*filter_argv_call)(void **),
        const struct halide_filter_metadata_t *filter_metadata,
        const char *const *extra_key_value_pairs) {
}
/*
 * HDRPlus Class -- Houses file I/O, defines pipeline attributes and calls
 * processes main stages of the pipeline.
 */
class HDRPlus {
    const Burst &burst;

public:
    const Compression c;
    const Gain g;

    HDRPlus(Burst& burst, const Compression c, const Gain g)
            : burst(burst), c(c), g(g)
    {
    }

    Halide::Runtime::Buffer<uint8_t> process() {
        const int width = burst.GetWidth();
        const int height = burst.GetHeight();

        Halide::Runtime::Buffer<uint8_t> output_img(3, width, height);

#ifdef _DEBUG
        std::cerr << "Black point: " << burst.GetBlackLevel() << std::endl;
        std::cerr << "White point: " << burst.GetWhiteLevel() << std::endl;
#endif

        const WhiteBalance wb = burst.GetWhiteBalance();
        std::cerr << "RGGB: " << wb.r << " " << wb.g0 << " " << wb.g1 << " " << wb.b
                  << std::endl;

        Halide::Runtime::Buffer<uint16_t> imgs = burst.ToBuffer();
        if (imgs.dimensions() != 3 || imgs.extent(2) < 2) {
            throw std::invalid_argument(
                    "The input of HDRPlus must be a 3-dimensional buffer with at least "
                    "two channels.");
        }

        const int cfa_pattern = static_cast<int>(burst.GetCfaPattern());
        auto ccm = burst.GetColorCorrectionMatrix();
        hdrplus_pipeline(imgs, burst.GetBlackLevel(), burst.GetWhiteLevel(), wb.r,
                         wb.g0, wb.g1, wb.b, cfa_pattern, ccm, c, g, output_img);

        // transpose to account for interleaved layout
        output_img.transpose(0, 1);
        output_img.transpose(1, 2);

        return output_img;
    }

#ifdef _DEBUG
    static bool save_png(const std::string &dir_path, const std::string &img_name,
                         const Halide::Runtime::Buffer<uint8_t> &img) {
        const std::string img_path = dir_path + "/" + img_name;

        const int stride_in_bytes = img.width() * img.channels();
        if (!stbi_write_png(img_path.c_str(), img.width(), img.height(),
                            img.channels(), img.data(), stride_in_bytes)) {
            std::cerr << "Unable to write output image '" << img_name << "'"
                      << std::endl;
            return false;
        }
        return true;
    }
#endif

};

bool doHdrPlus(const std::vector< std::vector<uint8_t> >& images, cv::Mat& mat)
{
    Compression c = 3.8f;
    Gain g = 1.1f;

    Burst burst(images);

    HDRPlus hdr_plus(burst, c, g);

    Halide::Runtime::Buffer<uint8_t> outputHdr = hdr_plus.process();

#ifdef _DEBUG
    HDRPlus::save_png("/sdcard/com.xypower.mpapp/tmp", "2.png", outputHdr);
#endif
    int width = outputHdr.width();
    int height = outputHdr.height();
    int channels = outputHdr.channels();
    int jch = 0;
    mat = cv::Mat::zeros(height, width, CV_8UC3);
    for (int i = 0; i < height; ++i)
    {
        jch = 0;
        for (int j = 0; j < width; ++j)
        {
            for (int n = 0; n < channels; ++n)
            {
                mat.at<uchar>(i, jch + n) = (uchar)outputHdr(j, i, n);
            }
            jch += channels;
        }
    }

    // if (!HDRPlus::save_png(dir_path, out_name, output)) {

    return true;
}

#if 0
int doHdrPlus(const std::string& dir_path, const std::string& out_name, const std::vector<std::string>& in_names) {


    Compression c = 3.8f;
    Gain g = 1.1f;

    Burst burst(dir_path, in_names);

    HDRPlus hdr_plus(burst, c, g);

    Halide::Runtime::Buffer<uint8_t> output = hdr_plus.process();

    if (!HDRPlus::save_png(dir_path, out_name, output)) {
        return EXIT_FAILURE;
    }

    return 0;
}
#endif