#pragma once

#include <vector>
#include <opencv2/opencv.hpp>

//
// INTERFACE
//

enum Histogram_matching_color_space
{
    HIST_MATCH_RGB,
    HIST_MATCH_HSV,
    HIST_MATCH_LAB
};


//
// Perform Histogram matching:
//
// The histogram of 'src' is adjusted so that it resembles the histogram of 'ref'
//
// If 'mask' is specified, the histogram matching will only be computed inside of the mask, but will be applied to the whole image!
//
static void histogram_matching(const cv::Mat& src, const cv::Mat& ref, cv::Mat& result, const cv::Mat& mask = cv::Mat(), Histogram_matching_color_space color_space = HIST_MATCH_LAB);




//
// Implementation
//


typedef std::vector<int>   Histogram;
typedef std::vector<float> Cumulative_distribution_function;
typedef std::vector<float> Matching_function;


//
// Opencv does not(!) store HSV and LAB in floating point range [0,1]
//
//  see https://docs.opencv.org/2.4/modules/imgproc/doc/miscellaneous_transformations.html#cvtcolor
//
// For histogram calculation however, we want values to be in this range. The two convert_* functions
// take care of mapping from and to the respective OpenCV representations
//
static cv::Vec3f convert_to_0_1(const cv::Vec3f val, const Histogram_matching_color_space color_space)
{
    switch (color_space)
    {
        // No conversion needed for RGB
        case HIST_MATCH_RGB: { return val; }

        // In OpenCV, Hue is stored from 0 to 360 (degrees)
        case HIST_MATCH_HSV: { return cv::Vec3f(val[0] / 360.f, val[1], val[2]); }

        case HIST_MATCH_LAB:
        {
            return cv::Vec3f( val[0]          / 100.f,  // In OpenCV:    0 < L < 100
                             (val[1] + 127.f) / 254.f,  // In OpenCV: -127 < a < 127
                             (val[2] + 127.f) / 254.f); // In OpenCV: -127 < b < 127
        }
    }

    return 0.0f;
}

static float convert_from_0_1(const float val, const int channel, const Histogram_matching_color_space color_space)
{
    switch (color_space)
    {
        // No conversion needed for RGB
        case HIST_MATCH_RGB: { return val; }

        case HIST_MATCH_HSV:
        {
            switch (channel)
            {
                // In OpenCV, Hue is stored from 0 to 360 (degrees)
                case 0: return val * 360.f;
                case 1: return val;
                case 2: return val;
            }
        }
        case HIST_MATCH_LAB:
        {
            switch (channel)
            {
                case 0: return val * 100.f;         // In OpenCV:    0 < L < 100
                case 1: return val * 254.f - 127.f; // In OpenCV: -127 < a < 127
                case 2: return val * 254.f - 127.f; // In OpenCV: -127 < b < 127

            }
        }
    }

    return 0.0f;
}

//
// Applies a matching function to the channel of the given image.
//
static void apply_match_function(cv::Mat& image, const cv::Mat& mask, const int channel, const Matching_function& matching_function, const Histogram_matching_color_space color_space)
{
    float num_bins = (float)matching_function.size();

    for (int row = 0; row < image.rows; ++row)
    {
        for (int col = 0; col < image.cols; ++col)
        {
            if (!mask.data || mask.at<uchar>(row, col) > 0)
            {
                // [0 to 1]
                cv::Vec3f current_val = convert_to_0_1(image.at<cv::Vec3f>(row, col), color_space);

                // [0 to 1]
                int   bin_index   = (int)(current_val[channel] * num_bins);
                float matched_val = matching_function[bin_index];

                // OpenCVs range
                float new_val = convert_from_0_1(matched_val, channel, color_space);

                // printf("[DEBUG] From colorspace  %f.  mapping  %f  ->  %f.   In colorspace  %f\n",
                //        image.at<cv::Vec3f>(row, col)[channel],
                //        current_val[channel],
                //        matched_val,
                //        new_val);

                image.at<cv::Vec3f>(row, col)[channel] = new_val;
            }
        }
    }
}

//
// Computes histogram in a given color space
//
// IMPORTANT(swenninger): image has to be of floating point type
//
static void compute_histogram(const cv::Mat& image, const cv::Mat& mask, const int channel, const int num_bins, const Histogram_matching_color_space color_space,
                              Histogram& histogram)
{
    //
    // Check image type
    //
    if (image.elemSize1() != sizeof(float))
    {
        printf("[ERROR] Histogram computiation is only implemented for float valued images\n");
        return;
    }

    histogram.clear();
    histogram.resize(num_bins, 0);

    for (int row = 0; row < image.rows; ++row)
    {
        for (int col = 0; col < image.cols; ++col)
        {
            if (mask.at<uchar>(row, col) == 255)
            {
                // [0 to 1]
                cv::Vec3f current_val = convert_to_0_1(image.at<cv::Vec3f>(row, col), color_space);

                // printf("Val (%f, %f, %f)\n", current_val[0], current_val[1], current_val[2]);

                int histogram_bin = (int)(current_val[channel] * num_bins);
                histogram[histogram_bin]++;
            }
        }
    }
}

static void compute_cdf(Histogram& hist, Cumulative_distribution_function& cdf)
{
    int num_bins = (int)hist.size();

    //
    // Clear and set to 0
    //
    cdf.clear();
    cdf.resize(hist.size(), 0.0f);

    //
    // Num pixels
    //
    float num_elems = 0;
    for (int count : hist)
        num_elems += (float)count;

    //
    // Compute cdf
    //
    float prob_bin_0 = ((float)hist[0]) / num_elems;

    // TODO(swenninger):  Is this correct or should this be 0.0f ???
    cdf[0] = prob_bin_0;

    for (int i = 1; i < num_bins; ++i)
    {
        float prob_bin_i = ((float)hist[i] / num_elems);

        cdf[i] = cdf[i - 1] + prob_bin_i;
    }
}

//
// Performs an average filter of the cdf signal
//
static void smooth_cdf(Cumulative_distribution_function& cdf)
{
    int num_bins = (int)cdf.size();

    Cumulative_distribution_function smoothed(cdf.size());

    //
    // TODO(swenninger): empirically determined (half) window size for average filter
    //
    int half_window_size = 10;
    int window_size = 2 * half_window_size + 1;

    for (int i = half_window_size; i < num_bins - half_window_size; ++i)
    {
        float val = 0.0;
        for (int j = -half_window_size; j <= half_window_size; ++j)
        {
            val += cdf[i + j];
        }
        smoothed[i] = val / window_size;
    }

    for (int i = half_window_size; i < num_bins - half_window_size; ++i)
        cdf[i] = smoothed[i];
}

//
// Mapping between src and ref
//
static void compute_matching_function(const Cumulative_distribution_function& src_cdf,
                                      const Cumulative_distribution_function& ref_cdf,
                                      Matching_function& matching_function)
{
    int num_bins = (int)src_cdf.size();
    float num_bins_f = (float)num_bins;

    matching_function.clear();
    matching_function.resize(src_cdf.size(), 0);

    float last_match = 0;
    for (int i = 0; i < num_bins; ++i)
    {
        float match_value = 0;
        for (int j = 0; j < num_bins; ++j)
        {
            if (ref_cdf[j] >= src_cdf[i])
            {
                match_value = (float)j / num_bins_f;
                last_match  = std::max(last_match, match_value);
                break;
            }
        }
        matching_function[i] = last_match;
    }

    //
    // NOTE(swenninger): This leads maps many color values to 0.0f, which leads to artefacts,
    //                   we smooth the mapping by linearly interpolating from 0.0f to the
    //                   first mapping entry above 0.0f.
    //
    int first_over_zero = 0;
    for (int i = 0; i < num_bins; ++i)
    {
        if (matching_function[i] > 0.0f)
        {
            first_over_zero = i;
            break;
        }
    }

    float first_over_zero_val = matching_function[first_over_zero];

    for (int i = 0; i < first_over_zero; ++i)
    {
        float t = (float)i / (float)first_over_zero;
        matching_function[i] = t * first_over_zero_val;
    }
}


//
// Debug function to visualize histograms
//
static void draw_histogram(const Histogram& histogram, cv::Mat& result_image)
{
    if (!result_image.data)
    {
        result_image = cv::Mat::zeros(512, histogram.size() * 2, CV_8UC1);
    }

    int width  = result_image.cols;
    int height = result_image.rows;

    int num_bins  = (int)histogram.size();
    int bin_width = (width / num_bins);

    // max
    int max = 0;
    for (int i = 0; i < num_bins; ++i) { if (histogram[i] > max) { max = histogram[i]; } }

    for (int i = 0; i < num_bins; ++i)
    {
        int hist_count  = histogram[i];
        int line_height = (int)height * (1.0 - ((double)hist_count / (double)max));

        cv::line(result_image,
                 cv::Point(i * bin_width, height),
                 cv::Point(i * bin_width, line_height),
                 cv::Scalar(255), 1, 8, 0);
    }
}

//
// NOTE(swenninger): Only works for 3-channel images with 32bit floats per channel!
//
static void filter_black(const cv::Mat& image, cv::Mat& mask)
{
    cv::Vec3f zero_vec3f(0.0f, 0.0f, 0.0f);
    for (int row = 0; row < image.rows; ++row)
        for (int col = 0; col < image.cols; ++col)
            if (image.at<cv::Vec3f>(row, col) == zero_vec3f)
                mask.at<uchar>(row, col) = 0;
}

//
// Histogram Matching algorithm
//
static void histogram_matching(const cv::Mat& src, const cv::Mat& ref, cv::Mat& result, const cv::Mat& mask, Histogram_matching_color_space color_space)
{
    int num_bins = 256;

    //
    // Convert to floating point
    //
    cv::Mat local_src;
    cv::Mat local_ref;
    src.convertTo(local_src, CV_32FC3, 1.0/255.0);
    ref.convertTo(local_ref, CV_32FC3, 1.0/255.0);

    //
    // Generate masks
    //
    cv::Mat src_mask = cv::Mat(local_src.size(), CV_8UC1, cv::Scalar(255));
    cv::Mat ref_mask = cv::Mat(local_ref.size(), CV_8UC1, cv::Scalar(255));

    if (mask.data && mask.channels() == 1)
    {
        src_mask = mask.clone();
        ref_mask = mask.clone();
    }
    else
    {
        if (mask.data)
        {
            printf("[WARNING] Histogram Matching was called with a mask with %d channels, expecting 1-channel image! Ignoring mask.\n", mask.channels());
        }

        filter_black(local_src, src_mask);
        filter_black(local_ref, ref_mask);
    }

    //
    // Color Space conversion
    //
    switch(color_space)
    {
        case HIST_MATCH_RGB:
        {
            break;
        }

        case HIST_MATCH_HSV:
        {
            cv::cvtColor(local_src, local_src, cv::COLOR_BGR2HSV);
            cv::cvtColor(local_ref, local_ref, cv::COLOR_BGR2HSV);
            break;
        }
        case HIST_MATCH_LAB:
        {
            cv::cvtColor(local_src, local_src, cv::COLOR_BGR2Lab);
            cv::cvtColor(local_ref, local_ref, cv::COLOR_BGR2Lab);
            break;
        }
    }

    result = local_src.clone();

    //
    // Which channels to match?
    //
    std::vector<int> channels_to_substitute;
    switch (color_space)
    {
        case HIST_MATCH_RGB: channels_to_substitute = {0, 1, 2};
        case HIST_MATCH_HSV: channels_to_substitute = {0, 1};    // Do not change V(alue)
        case HIST_MATCH_LAB: channels_to_substitute = {1, 2};    // Do not change L(uminosity)
    }

    //
    // Per-channel matching
    //
    for (auto channel : channels_to_substitute)
    {
        Histogram src_hist, ref_hist;
        Cumulative_distribution_function src_cdf, ref_cdf;
        Matching_function match;

        compute_histogram(local_src, src_mask, channel, num_bins, color_space, src_hist);
        compute_histogram(local_ref, ref_mask, channel, num_bins, color_space, ref_hist);

        compute_cdf(src_hist,  src_cdf);
        compute_cdf(ref_hist, ref_cdf);

        smooth_cdf(src_cdf);
        smooth_cdf(ref_cdf);

        compute_matching_function(src_cdf, ref_cdf, match);

        if (mask.data)
        {
            // If the user specified a mask, where the mapping function should be calculated,
            // we still want to apply it to all of the image!
            apply_match_function(result, cv::Mat(), channel, match, color_space);
        }
        else
        {
            apply_match_function(result, src_mask, channel, match, color_space);
        }

        //
        // DEBUG
        //

        // Histogram res_hist;
        // compute_histogram(result, src_mask, channel, num_bins, color_space, res_hist);

#if 0
        //
        // TODO(swenninger) remove. This writes output for analyzing histograms, cdfs and matching function (e.g via gnuplot)
        //
        std::ofstream out;
        out = std::ofstream("src_hist_" + std::to_string(channel) + ".txt");
        for (size_t i = 0; i < match.size(); ++i) out << (float)i / (float)match.size() << " " << src_hist[i] << std::endl;

        out = std::ofstream("ref_hist_" + std::to_string(channel) + ".txt");
        for (size_t i = 0; i < match.size(); ++i) out << (float)i / (float)match.size() << " " << ref_hist[i] << std::endl;

        out = std::ofstream("src_cdf_" + std::to_string(channel) + ".txt");
        for (size_t i = 0; i < match.size(); ++i) out << (float)i / (float)match.size() << " " << src_cdf[i] << std::endl;

        out = std::ofstream("ref_cdf_" + std::to_string(channel) + ".txt");
        for (size_t i = 0; i < match.size(); ++i) out << (float)i / (float)match.size() << " " << ref_cdf[i] << std::endl;

        out = std::ofstream("match_" + std::to_string(channel) + ".txt");
        for (size_t i = 0; i < match.size(); ++i) {out << (float)i / (float)match.size() << " " << match[i] << std::endl; }

        out = std::ofstream("res_hist_" + std::to_string(channel) + ".txt");
        for (size_t i = 0; i < match.size(); ++i) out << (float)i / (float)match.size() << " " << res_hist[i] << std::endl;
#endif

        // cv::Mat vis_src_hist;
        // cv::Mat vis_ref_hist;
        // cv::Mat vis_res_hist;
        // draw_histogram(src_hist, vis_src_hist);
        // draw_histogram(ref_hist, vis_ref_hist);
        // draw_histogram(res_hist, vis_res_hist);
        // cv::imshow("src_hist", vis_src_hist);
        // cv::imshow("ref_hist", vis_ref_hist);
        // cv::imshow("res_hist", vis_res_hist);
        // cv::waitKey(0);
    }

    //
    // Return back to original color space
    //
    switch(color_space)
    {
        case HIST_MATCH_HSV: { cv::cvtColor(result, result, cv::COLOR_HSV2BGR); break; }
        case HIST_MATCH_LAB: { cv::cvtColor(result, result, cv::COLOR_Lab2BGR); break; }
    }

    result.convertTo(result, CV_8UC3, 255);
}
