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

#define max(a, b) ((a) > (b) ? (a) : (b))
#define SIZE_OF_TABLE 256 * 256 * 256

CImageApplyColorCastCorrect::CImageApplyColorCastCorrect()
	: m_table(new uint[SIZE_OF_TABLE])
{
    std::vector<double> points_x, points_y;
    points_x = { 0, 80, 175, 255 };
    points_y = { 12, 85, 175, 270 };
    createTable(points_x, points_y);
}

CImageApplyColorCastCorrect::CImageApplyColorCastCorrect(const std::vector<double>& points_x, const std::vector<double>& points_y)
	: m_table(new uint[256 * 256 * 256])
{
    createTable(points_x, points_y);
}

CImageApplyColorCastCorrect::CImageApplyColorCastCorrect(const std::string& fileName)
    : m_table(new uint[256 * 256 * 256])
{
    std::fstream file(fileName, std::ios::in | std::ios::binary);
    if (file)
        file.read(reinterpret_cast<char*>(m_table), SIZE_OF_TABLE * sizeof(uint));
    file.close();
}

CImageApplyColorCastCorrect::~CImageApplyColorCastCorrect(void)
{
	delete[] m_table;
}

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

    cv::Mat bgra;
    cv::cvtColor(pDib, bgra, cv::COLOR_BGR2BGRA);
    uint* ptr = bgra.ptr<uint>();
    int rows = bgra.rows, cols = bgra.cols;
    for (int i = 0; i < rows; i++)
    {
        ptr = reinterpret_cast<uint*>(bgra.ptr(i));
        for (int j = 0; j < cols; j++)
            ptr[j] = m_table[ptr[j] & 0x00ffffff];
    }
    cv::cvtColor(bgra, pDib, cv::COLOR_BGRA2BGR);
    bgra.release();
}

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), SIZE_OF_TABLE * sizeof(uint));
    file.close();
}

std::vector<double> CImageApplyColorCastCorrect::caculate(const std::vector<double>& points_x, const std::vector<double>& points_y)
{
    int MaxElement = points_x.size() - 1;
    //计算常数f
    double f = points_y[0];
    //求解
    int n, m;
    //double a[MaxElement][MaxElement+1];
    std::vector<std::vector<double>> a;
    //a.resize(MaxElement);
    for (int i = 0; i < MaxElement; i++)
    {
        std::vector<double> b;
        b.resize(MaxElement + 1);
        a.push_back(b);
    }

    for (int i = 0; i < MaxElement; i++)
    {
        for (int j = 0; j < MaxElement; j++)
            a[i][j] = cv::pow(points_x[i + 1], MaxElement - j);
        a[i][MaxElement] = points_y[i + 1] - f;
    }

    int i, j;
    n = MaxElement;

    for (j = 0; j < n; j++)
    {
        double max = 0;
        double imax = 0;
        for (i = j; i < n; i++)
            if (imax < cv::abs(a[i][j]))
            {
                imax = cv::abs(a[i][j]);
                max = a[i][j];//得到各行中所在列最大元素
                m = i;
            }

        if (cv::abs(a[j][j]) != max)
        {
            double b = 0;
            for (int k = j; k < n + 1; k++)
            {
                b = a[j][k];
                a[j][k] = a[m][k];
                a[m][k] = b;
            }
        }

        for (int r = j; r < n + 1; r++)
            a[j][r] = a[j][r] / max;//让该行的所在列除以所在列的第一个元素，目的是让首元素为1

        for (i = j + 1; i < n; i++)
        {
            double c = a[i][j];
            if (c == 0.0) continue;
            for (int s = j; s < n + 1; s++)
                a[i][s] = a[i][s] - a[j][s] * c;//前后行数相减，使下一行或者上一行的首元素为0
        }
    }

    for (i = n - 2; i >= 0; i--)
        for (j = i + 1; j < n; j++)
            a[i][n] = a[i][n] - a[j][n] * a[i][j];

    std::vector<double> result;
    for (int k = 0; k < n; k++)
        result.push_back(a[k][n]);
    result.push_back(f);
    return result;
}

void CImageApplyColorCastCorrect::createTable(const std::vector<double>& points_x, const std::vector<double>& points_y)
{
    cv::Mat mat_rgbTable(256 * 256, 256, CV_8UC3);
    uchar* ptr_mat_rgbTable = mat_rgbTable.data;
    for (size_t r = 0; r < 256; r++)
        for (size_t g = 0; g < 256; g++)
            for (size_t b = 0; b < 256; b++, ptr_mat_rgbTable += 3)
            {
                ptr_mat_rgbTable[0] = b;
                ptr_mat_rgbTable[1] = g;
                ptr_mat_rgbTable[2] = r;
            }

    cv::cvtColor(mat_rgbTable, mat_rgbTable, cv::COLOR_BGR2HSV_FULL);

    uchar table_data[256];
    cv::Mat tableLUT(256, 1, CV_8UC1, table_data);
    std::vector<double> coefficient;

    coefficient = caculate(points_x, points_y);
    for (int j = 0; j < 256; j++)
        table_data[j] = static_cast<int>(max(0, coefficient[0] * j * j * j + coefficient[1] * j * j + coefficient[2] * j + coefficient[3]));

    cv::Mat hsv_ms[3];
    cv::split(mat_rgbTable, hsv_ms);
    cv::LUT(hsv_ms[0], tableLUT, hsv_ms[0]);
    cv::merge(hsv_ms, 3, mat_rgbTable);

    cv::cvtColor(mat_rgbTable, mat_rgbTable, cv::COLOR_HSV2BGR_FULL);

    cv::Mat mat_bgr32(256, 256 * 256, CV_8UC4);
    cv::cvtColor(mat_rgbTable, mat_bgr32, cv::COLOR_BGR2BGRA);

    memcpy(m_table, mat_bgr32.data, mat_bgr32.total() * sizeof(uint));
}