#include "ImageApplyColorCastCorrect.h"
#include <fstream>
#include <iostream>

#define JIANG_BAOHE

#define max(a, b) ((a) > (b) ? (a) : (b))
#define min(a, b) ((a) < (b) ? (a) : (b))
constexpr auto SIZE_OF_TABLE = 256;

CImageApplyColorCastCorrect::CImageApplyColorCastCorrect(PreScheme ps)
    : m_table_h(new uchar[SIZE_OF_TABLE])
    , m_table_l(new uchar[SIZE_OF_TABLE])
    , m_table_s(new uchar[SIZE_OF_TABLE])
{
    std::vector<double> points_x, points_y;
    
    switch (ps)
    {
    case CImageApplyColorCastCorrect::G200:
        points_x = { 2,  9, 15, 18, 22, 34, 50, 63,  78,  92, 103, 114, 126, 137, 140, 147, 154, 163, 176, 200, 225, 248, 251, 254 };
        points_y = { 11, 17, 24, 28, 32, 41, 56, 81, 106, 115, 119, 124, 131, 138, 141, 145, 153, 166, 186, 201, 217, 251, 259, 264 };
        break;
    case CImageApplyColorCastCorrect::G300:
        points_x = { 1,  7, 12, 18, 33, 52, 68, 81,  91, 100, 111, 125, 138, 142, 147, 153, 161, 172, 198, 230, 248, 250, 252, 253 };               
        points_y = { 26, 31, 33, 36, 40, 44, 56, 92, 104, 114, 126, 135, 141, 143, 146, 151, 169, 198, 218, 227, 252, 266, 272, 276 };
        break;
    case CImageApplyColorCastCorrect::G300_D8:
        points_x = { 0,  2,  5,  9, 15, 20, 26, 39, 54, 70, 83,  94, 102, 113, 128, 139, 146, 151, 156, 164, 177, 199, 232, 253 };
        points_y = { 11, 17, 21, 26, 31, 33, 36, 40, 44, 56, 92, 104, 114, 126, 135, 141, 143, 146, 151, 169, 198, 218, 227, 252 };
        break;
    case CImageApplyColorCastCorrect::G400_402:
        points_x = { 2,  8, 15, 20, 27, 41, 50, 67, 80,  91,  99, 109, 124, 138, 145, 150, 157, 164, 177, 202, 232, 250, 252, 255 };
        points_y = { 21, 26, 31, 33, 36, 40, 44, 56, 92, 104, 114, 126, 135, 141, 143, 146, 151, 169, 198, 218, 227, 252, 266, 272 };
        break;
    case CImageApplyColorCastCorrect::G400_3288:
        points_x = { 2,  3,  5,  7, 11, 16, 19, 25, 40, 57, 72, 84,  95, 101, 111, 126, 144, 148, 152, 157, 165, 179, 203, 234 };
        points_y = { -3, 11, 17, 21, 26, 31, 33, 36, 40, 44, 56, 92, 104, 114, 126, 135, 141, 143, 146, 151, 169, 198, 218, 227 };
        break;
    case CImageApplyColorCastCorrect::Android302:
        points_x = { 0, 6,  11, 19, 37, 48, 66, 79,  89,  97, 108, 122, 138, 143, 147, 150, 158, 170, 199, 233, 248, 249, 250, 252 };
        points_y = { 26, 31, 33, 36, 40, 44, 56, 92, 104, 114, 126, 135, 141, 143, 146, 151, 169, 198, 218, 227, 252, 266, 272, 276 };
        break;
    case CImageApplyColorCastCorrect::G300_7010:
        points_x = { 1, 4, 6, 9, 13, 19, 23, 29, 40, 55, 72, 86, 97, 106, 117, 133, 143, 149, 155, 160, 170, 182, 206, 233 };
        points_y = { 252,  11,  17,  21, 26, 31, 33, 36, 40, 44, 56,  92, 104, 114, 126, 135, 141, 143, 146, 151, 169, 198, 218, 227 };
        break;
    default:
        points_x = { 0, 255 };
        points_y = { 0, 255 };
        break;
    }

    createTable_h(points_x, points_y);

    points_x = { 0, 180, 255 };
    points_y = { 0, 175, 220 };
    createTable_l(points_x, points_y);

    points_x = { 0, 230, 255 };
    points_y = { 0, 210, 255 };
    createTable_s(points_x, points_y);
}

CImageApplyColorCastCorrect::CImageApplyColorCastCorrect(const std::vector<double>& h_x, const std::vector<double>& h_y, 
    const std::vector<double>& l_x, const std::vector<double>& l_y, 
    const std::vector<double>& s_x, const std::vector<double>& s_y)
    : m_table_h(new uchar[SIZE_OF_TABLE])
    , m_table_l(new uchar[SIZE_OF_TABLE])
    , m_table_s(new uchar[SIZE_OF_TABLE])
{
    if (!h_x.empty() && !h_y.empty())
        createTable_h(h_x, h_y);
    if (!l_x.empty() && !l_y.empty())
        createTable_l(l_x, l_y);
    if (!s_x.empty() && !s_y.empty())
        createTable_s(s_x, s_y);
}

CImageApplyColorCastCorrect::CImageApplyColorCastCorrect(const std::vector<double>& points_x, const std::vector<double>& points_y)
	: m_table_h(new uchar[SIZE_OF_TABLE])
    , m_table_l(new uchar[SIZE_OF_TABLE])
    , m_table_s(new uchar[SIZE_OF_TABLE])
{
    createTable_h(points_x, points_y);

    for (size_t i = 0; i < SIZE_OF_TABLE; i++)
        m_table_l[i] = m_table_s[i] = i;
}

CImageApplyColorCastCorrect::CImageApplyColorCastCorrect(const std::string& fileName)
    : m_table_h(new uchar[SIZE_OF_TABLE])
    , m_table_l(new uchar[SIZE_OF_TABLE])
    , m_table_s(new uchar[SIZE_OF_TABLE])
{
    std::fstream file(fileName, std::ios::in | std::ios::binary);
    if (file)
        file.read(reinterpret_cast<char*>(m_table_h), SIZE_OF_TABLE);
    file.close();

    for (size_t i = 0; i < SIZE_OF_TABLE; i++)
        m_table_l[i] = m_table_s[i] = i;
}

CImageApplyColorCastCorrect::CImageApplyColorCastCorrect(const uchar* table_h)
    : m_table_h(new uchar[SIZE_OF_TABLE])
    , m_table_l(new uchar[SIZE_OF_TABLE])
    , m_table_s(new uchar[SIZE_OF_TABLE])
{
    memcpy(m_table_h, table_h, SIZE_OF_TABLE);
    for (size_t i = 0; i < SIZE_OF_TABLE; i++)
        m_table_l[i] = m_table_s[i] = i;
}

CImageApplyColorCastCorrect::~CImageApplyColorCastCorrect(void)
{
	delete[] m_table_h;
    delete[] m_table_l;
    delete[] m_table_s;
}

void CImageApplyColorCastCorrect::apply(cv::Mat& pDib, int side)
{
    if (pDib.channels() != 3)
        return;
    cv::Mat hsv;
    cv::cvtColor(pDib, hsv, cv::COLOR_BGR2HLS_FULL);

    cv::Mat hsv_mv[3];
    cv::split(hsv, hsv_mv);

    //cv::imwrite("hsv_0.jpg", hsv_mv[0]);
    cv::Mat lut(256, 1, CV_8UC1, m_table_h);
    cv::LUT(hsv_mv[0], lut, hsv_mv[0]);

#ifdef JIANG_BAOHE
    //hsv_mv[2] -= 20;
    //cv::imwrite("hsv_0.jpg", hsv_mv[0]);
    //cv::imwrite("hsv_1.jpg", hsv_mv[1]);
    //cv::imwrite("hsv_2.jpg", hsv_mv[2]);
    cv::Mat lut2(256, 1, CV_8UC1, m_table_l);
    cv::LUT(hsv_mv[2], lut2, hsv_mv[2]);

    //hsv_mv[1] -= 20;
    //cv::Mat lut3(256, 1, CV_8UC1, m_table3);
    //cv::LUT(hsv_mv[1], lut3, hsv_mv[1]);
#endif
    cv::merge(hsv_mv, 3, pDib);

    cv::cvtColor(pDib, pDib, cv::COLOR_HLS2BGR_FULL);
}

void CImageApplyColorCastCorrect::apply(std::vector<cv::Mat>& mats, bool isTwoSide)
{
    (void)isTwoSide;
    int i = 0;
    for (cv::Mat& var : mats) {
        if (i != 0 && isTwoSide == false)
            break;
        if (!var.empty())
            apply(var, 0);
        i++;
    }
}

void CImageApplyColorCastCorrect::exportTableData(const std::string& fileName)
{
    std::fstream file(fileName, std::ios::out | std::ios::binary);
    if (file)
        file.write(reinterpret_cast<char*>(m_table_h), SIZE_OF_TABLE);
    file.close();
}

void CImageApplyColorCastCorrect::createTable_h(const std::vector<double>& points_x, const std::vector<double>& points_y)
{
    int table_temp[256]{};

    for (size_t i = 0; i < points_x.size(); i++)
    {
        int current_index = static_cast<int>(points_x[i]);
        if (current_index == 255)
            current_index = 0;
        int next_index = static_cast<int>(points_x[(i + 1) % points_x.size()]);

        double low = points_y[i];
        double up = points_y[(i + 1) % points_y.size()];
        if (low == 255)
            low = 0;
        if (up < low)
            up += 255;

        if (next_index < current_index)
            next_index += 256;

        int length = next_index - current_index + 1;
        double step = (up - low) / length;

        for (int j = 0; j < length; j++)
        {
            int temp = (j + current_index) % 256;
            table_temp[temp] = step * j + low;
        }

        for (size_t j = 0; j < 256; j++)
            if (table_temp[j] > 255)
                m_table_h[j] = table_temp[j] - 255;
            else
                m_table_h[j] = table_temp[j];
    }
}

void CImageApplyColorCastCorrect::createTable_l(const std::vector<double>& points_x, const std::vector<double>& points_y)
{
    memset(m_table_l, 255, 256);
    memset(m_table_l, 0, 127);
    int s = 1;
    for (int i = 0; i < points_x.size() - s; i += s)
    {
        double low = points_y[i];
        double up = points_y[i + s];

        int start = points_x[i];
        int end = points_x[i + s];

        int length = end - start;
        double step = (up - low) / length;
        for (int j = 0; j < length; j++)
            m_table_l[start + j] = max(0.0, min(255.0, low + j * step));
    }
}

void CImageApplyColorCastCorrect::createTable_s(const std::vector<double>& points_x, const std::vector<double>& points_y)
{
    memset(m_table_s, 255, 256);
    memset(m_table_s, 0, 127);
    int s = 1;
    for (int i = 0; i < points_x.size() - s; i += s)
    {
        double low = points_y[i];
        double up = points_y[i + s];

        int start = points_x[i];
        int end = points_x[i + s];

        int length = end - start;
        double step = (up - low) / length;
        for (int j = 0; j < length; j++)
            m_table_s[start + j] = max(0.0, min(255.0, low + j * step));
    }
}