#include "correctProcess.h"
#include <opencv2/core/types_c.h>
#include <opencv2/core/core_c.h>
#include <opencv2/imgproc/imgproc_c.h>
#include <iostream>
#include <fstream>

bool getMeanColors(const cv::Mat& image, std::vector<std::vector<uchar>>& colors)
{
    std::vector<std::vector<int>> sections;
    cv::Mat resizeMat(image.rows, 2, CV_8UC1);
    cv::resize(image(cv::Rect(0, 0, 10, image.rows)), resizeMat(cv::Rect(0, 0, 1, image.rows)), cv::Size(1, image.rows));
    cv::resize(image(cv::Rect(image.cols - 11, 0, 10, image.rows)), resizeMat(cv::Rect(1, 0, 1, image.rows)), cv::Size(1, image.rows));

    for (int i = 0; i < resizeMat.cols; i++)
    {
        cv::Mat col = resizeMat(cv::Rect(i, 0, 1, resizeMat.rows)).clone();
        std::vector<int> section;
        uchar* ptr = col.data;
        int thre = 55;
        for (int j = 2, rows = col.rows - 2; j < rows; j++)
        {
            if (abs(ptr[j - 2] - ptr[j + 2]) > thre)
            {
                if (section.size() > 1)
                {
                    if (j - *(section.end() - 1) < 15)
                        continue;
                }
                else if (section.size() == 1)
                    if (j - section[0] < 50)
                        continue;
                section.push_back(j);
                j += 5;
            }
        }

        if (section.size() < 53)
            return false;

        section.erase(section.begin());
        sections.push_back(section);
    }
    std::vector<double> steps;
    for (int i = 0, length = 52; i < length; i++)
        steps.push_back(static_cast<double>(sections[1][i] - sections[0][i]) / (image.cols - 1));

    for (int i = 0; i < image.cols; i++)
    {
        cv::Mat col = image(cv::Rect(i, 0, 1, image.rows)).clone();
        std::vector<uchar> section_color;
        for (size_t j = 0; j < 52; j += 2)
        {
            int top = sections[0][j] + i * steps[j] + 5;
            int bottom = sections[0][j + 1] + i * steps[j + 1] - 5;
            if (bottom <= top)
                return false;
            section_color.push_back(cv::mean(col(cv::Rect(0, top, 1, bottom - top + 1)))[0]);
        }

        for (size_t j = 1; j < section_color.size() - 1; j++)
            if (section_color[j] < section_color[j - 1] || section_color[j] > section_color[j + 1])
                section_color[j] = (int)(section_color[j - 1] + section_color[j + 1]) / 2;

        colors.push_back(section_color);
    }
    return true;
}

bool calTable(const cv::Mat& image, uchar* lut_data, int blackLevel, char* data)
{
    std::vector<std::vector<uchar>> colors;
    if (!getMeanColors(image, colors))
        return false;

    int level = blackLevel;

#if 1
    int index = 0;
    for (size_t i = 0; i < colors.size(); i++)
        for (size_t j = 0, length = colors[i].size(); j < length; j++)
        {
            data[index] = colors[i][j];
            index++;
        }
#endif

    for (int i = 0; i < colors.size(); i++)
        for (int j = 0; j < level; j++)
            colors[i].erase(colors[i].begin());

    double ss = 250.0 / (colors[0].size() - 1);

    for (int i = 0, length = image.cols; i < length; i++)
    {
        uchar* ptr = lut_data + 256 * i;
        memset(ptr, 255, 256);
        memset(ptr, 0, 100);

        for (int y = 1; y < colors[i].size(); y++)
        {
            int index_0 = colors[i][y - 1];
            int index_1 = colors[i][y];

            double value_0 = (y - 1) * ss;
            double value_1 = cv::min(y * ss, 255.0);
            double step = (double)(value_1 - value_0) / (double)(index_1 - index_0);
            for (int x = index_0; x < index_1; x++)
                ptr[x] = cv::min((int)(value_0 + (x - index_0) * step), 255);
        }
    }

    return true;
}

cv::Mat createTable(const cv::Mat& image, int blackLevel, char* data)
{
    cv::Mat grey(image.rows, image.cols * image.channels(), CV_8UC1, image.data);
    cv::Mat lut_data(grey.cols, 256, CV_8UC1);

    if (!calTable(grey, lut_data.data, blackLevel, data))
        return cv::Mat();

    const int channels = (lut_data.rows % 408 == 0) ? 408 : 432;

    for (size_t i = 0, block = lut_data.rows / channels; i < block; i++)
    {
        cv::Mat lutROI = lut_data(cv::Rect(0, i * channels, 256, channels));
        cv::Mat tran;
        cv::transpose(lutROI, tran);
        memcpy(lutROI.data, tran.data, tran.total());
    }

    return lut_data;
}

int mean(const cv::Mat& image, const cv::Mat& mask)
{
    cv::Mat mask1, mask2;
    cv::threshold(image, mask1, 127, 255, cv::THRESH_BINARY);
    mask2 = (~mask1) & mask;
    mask1 &= mask;

    double num1 = cv::sum(mask1)[0] / 255;
    double num2 = cv::sum(mask2)[0] / 255;

    double sum1 = cv::sum(mask1 & image)[0];
    double sum2 = cv::sum(mask2 & image)[0];

    if (num1 == 0 && num2 == 0)
        return 0;
    else if (num1 == 0 && num2 != 0)
        return sum2 / num2;
    else if (num1 != 0 && num2 == 0)
        return sum1 / num1;
    else if (abs(sum1 / num1 - sum2 / sum2) < 200)
        return (sum1 + sum2) / (num1 + num2);
    else
        return (sum1 + sum2 + num2 * 255) / (num1 + num2);
}

void myFindContours(const cv::Mat& src, std::vector<std::vector<cv::Point>>& contours, std::vector<cv::Vec4i>& hierarchy, int retr, int method, cv::Point offset)
{
    CvMat c_image = cvMat(src);

    cv::MemStorage storage(cvCreateMemStorage());
    CvSeq* _ccontours = 0;

    cvFindContours(&c_image, storage, &_ccontours, sizeof(CvContour), retr, method, cvPoint(offset.x, offset.y));

    if (!_ccontours)
    {
        contours.clear();
        return;
    }
    cv::Seq<CvSeq*> all_contours(cvTreeToNodeSeq(_ccontours, sizeof(CvSeq), storage));
    int total = (int)all_contours.size();
    contours.resize(total);

    cv::SeqIterator<CvSeq*> it = all_contours.begin();
    for (int i = 0; i < total; i++, ++it)
    {
        CvSeq* c = *it;
        ((CvContour*)c)->color = (int)i;
        int count = (int)c->total;
        int* data = new int[count * 2];
        cvCvtSeqToArray(c, data);
        for (int j = 0; j < count; j++)
        {
            contours[i].push_back(cv::Point(data[j * 2], data[j * 2 + 1]));
        }
        delete[] data;
    }

    hierarchy.resize(total);
    it = all_contours.begin();
    for (int i = 0; i < total; i++, ++it)
    {
        CvSeq* c = *it;
        int h_next = c->h_next ? ((CvContour*)c->h_next)->color : -1;
        int h_prev = c->h_prev ? ((CvContour*)c->h_prev)->color : -1;
        int v_next = c->v_next ? ((CvContour*)c->v_next)->color : -1;
        int v_prev = c->v_prev ? ((CvContour*)c->v_prev)->color : -1;
        hierarchy[i] = cv::Vec4i(h_next, h_prev, v_next, v_prev);
    }
    storage.release();
}

void myFillPolys(cv::Mat& image, const std::vector<std::vector<cv::Point>>& contours, const cv::Scalar& color)
{
    if (contours.empty()) return;

    size_t count = contours.size();
    cv::Point** pointss = new cv::Point * [count];
    int* npts = new int[count];

    for (size_t i = 0; i < count; i++)
    {
        size_t length = contours[i].size();
        npts[i] = length;
        pointss[i] = new cv::Point[length];
        for (size_t j = 0; j < length; j++)
            pointss[i][j] = contours[i][j];
    }
    cv::fillPoly(image, const_cast<const cv::Point**>(pointss), npts, count, color);

    for (size_t i = 0; i < count; i++)
        delete[] pointss[i];

    delete[] pointss;
    delete[] npts;
}

void createTable(const std::vector<double>& points_x, const std::vector<uchar>& points_y, uchar* table)
{
    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;
        if (current_index < 0)
            current_index += 255;
        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++)
            table_temp[(j + current_index) % 256] = step * j + low;

        for (size_t j = 0; j < 256; j++)
            table[j] = table_temp[j] % 255;
    }
}

cv::Mat createColorCastTable(const cv::Mat& image, const std::vector<uchar>& points_y)
{
    cv::Mat resizeMat;
    cv::resize(image, resizeMat, cv::Size(800, 800));

    cv::Mat hsv, mv[3];
    cv::cvtColor(resizeMat, hsv, cv::COLOR_BGR2HSV_FULL);
    cv::split(hsv, mv);

    cv::threshold(mv[1], mv[1], 80, 255, cv::THRESH_BINARY);
    cv::threshold(mv[2], mv[2], 30, 255, cv::THRESH_BINARY);

    cv::Mat mask = mv[1] & mv[2];
    cv::erode(mask, mask, cv::Mat());

    std::vector<std::vector<cv::Point>> contours;
    std::vector<cv::Vec4i> hierarchy;

    myFindContours(mask, contours, hierarchy, 0, 2, cv::Point(-1, -1));

    for (std::vector<std::vector<cv::Point>>::iterator iter = contours.begin(); iter != contours.end();)
        if (cv::contourArea(*iter) < 400)
            iter = contours.erase(iter);
        else
            iter++;

    if (contours.size() < 12)
        return cv::Mat();

    for (size_t i = 0; i < contours.size() - 1; i++)
        for (size_t j = i + 1; j < contours.size(); j++)
            if (contours[j][0].x < contours[i][0].x)
                std::swap(contours[j], contours[i]);

    std::vector<double> points_x;
    std::vector<std::vector<cv::Point>> contours_temp;
    for (size_t i = 0; i < 12; i++)
    {
        contours_temp.clear();
        contours_temp.push_back(contours[i]);
        mask = cv::Mat::zeros(resizeMat.size(), CV_8UC1);
        myFillPolys(mask, contours_temp, cv::Scalar::all(255));
        int mean_value = mean(mv[0], mask);
        if (mean_value >= 255)
            mean_value -= 255;
        points_x.push_back(mean_value);
    }

    if (points_x[0] + points_x[1] + points_x[2] > points_x[8] + points_x[9] + points_x[10])
        std::reverse(points_x.begin(), points_x.end());

    cv::Mat table(256, 1, CV_8UC1);
    createTable(points_x, points_y, table.data);
    return table;
}

void correctLUT(cv::Mat& image, const cv::Mat& lut)
{
    cv::Mat image_temp(image.rows, image.step / lut.channels(), CV_8UC(lut.channels()), image.data);
    for (size_t i = 0; i < image_temp.cols; i++)
        cv::LUT(image_temp(cv::Rect(i, 0, 1, image_temp.rows)), lut(cv::Rect(0, i, 256, 1)), image_temp(cv::Rect(i, 0, 1, image_temp.rows)));
}

std::vector<cv::Mat> correctProcess(cv::Mat& image, int blackLevel, bool colorCastEnable, const std::vector<uchar>& referValues, int dpi, const std::string& path)
{
    std::vector<cv::Mat> ret;
    if (image.empty())
        return ret;

    char* data = new char[image.step * 26];
    
    cv::Mat lut1 = createTable(image(cv::Rect(0, 0, image.cols / 2, image.rows)).clone(), blackLevel, data);

    if (lut1.empty())
    {
        delete[] data;
        return ret;
    }
	
    cv::Mat lut2 = createTable(image(cv::Rect(image.cols / 2, 0, image.cols / 2, image.rows)).clone(), blackLevel, data + image.step * 13);

    if (lut2.empty())
    {
        delete[] data;
        return ret;
    }

    std::ofstream file;
    file.open(path + std::to_string(dpi) + (image.channels() == 3 ? "c" : "g") + ".dat", std::ios::binary | std::ios::out);
    file.write(data, image.step * 26);
    file.close();
    delete[] data;

	cv::Mat lut(lut1.rows + lut2.rows, 256, CV_8UC1);
	lut1.copyTo(lut(cv::Rect(0, 0, 256, lut1.rows)));
    lut2.copyTo(lut(cv::Rect(0, lut1.rows, 256, lut2.rows)));
    ret.push_back(lut);

    if (!blackLevel || image.channels() == 1)
        return ret;

    int channel = (lut.rows % 408 == 0) ? 408 : 432;
    cv::Mat lutMat(lut.rows / channel, 256, CV_8UC(channel), lut.data);
    correctLUT(image, lutMat);

    cv::Mat colorCastTable = createColorCastTable(image(cv::Rect(0, 0, image.cols / 2, image.rows)), referValues);
    ret.push_back(colorCastTable);

    if (!colorCastTable.empty())
    {
        file.open(path + "cc" + std::to_string(dpi) + "f.dat", std::ios::binary | std::ios::out);
        file.write((char*)colorCastTable.data, 256);
        file.close();
    }

    colorCastTable = createColorCastTable(image(cv::Rect(image.cols / 2, 0, image.cols / 2, image.rows)), referValues);
    ret.push_back(colorCastTable);

    if (!colorCastTable.empty())
    {
        file.open(path + "cc" + std::to_string(dpi) + "b.dat", std::ios::binary | std::ios::out);
        file.write((char*)colorCastTable.data, 256);
        file.close();
    }

    return ret;
}

std::vector<cv::Mat> correctProcess(cv::Mat& image1, cv::Mat& image2, int blackLevel, bool colorCastEnable, const std::vector<uchar>& referValues, int dpi, const std::string& path)
{
    std::vector<cv::Mat> ret;

    char* data = new char[(image1.step + image2.step) * 26];
    cv::Mat lut1 = createTable(image1, blackLevel, data);

    if (lut1.empty())
        return ret;
	
    cv::Mat lut2 = createTable(image2, blackLevel, data);

    if (lut2.empty())
        return ret;

    std::ofstream file;
    file.open(path + std::to_string(dpi) + (image1.channels() == 3 ? "c" : "g") + ".dat", std::ios::binary | std::ios::out);
    file.write(data, (image1.step + image2.step) * 26);
    file.close();
    delete[] data;

	cv::Mat lut(lut1.rows + lut2.rows, 256, CV_8UC1);
	lut1.copyTo(lut(cv::Rect(0, 0, 256, lut1.rows)));
    lut2.copyTo(lut(cv::Rect(0, lut1.rows, 256, lut2.rows)));
    ret.push_back(lut);

    if (!blackLevel || image1.channels() == 1 || image2.channels() == 1)
        return ret;

	cv::Mat image(cv::max(image1.rows, image2.rows), image1.cols + image2.cols, CV_8UC3);
	
	image1.copyTo(image(cv::Rect(0, 0, image1.cols, image1.rows)));
    image2.copyTo(image(cv::Rect(image1.cols, 0, image2.cols, image2.rows)));
	
    int channel = (lut.rows % 408 == 0) ? 408 : 432;
    cv::Mat lutMat(lut.rows / channel, 256, CV_8UC(channel), lut.data);
    correctLUT(image, lutMat);

    cv::Mat colorCastTable = createColorCastTable(image(cv::Rect(0, 0, image.cols / 2, image.rows)), referValues);
    ret.push_back(colorCastTable);

    if (!colorCastTable.empty())
    {
        file.open(path + "cc" + std::to_string(dpi) + "f.dat", std::ios::binary | std::ios::out);
        file.write((char*)colorCastTable.data, 256);
        file.close();
    }

    colorCastTable = createColorCastTable(image(cv::Rect(image.cols / 2, 0, image.cols / 2, image.rows)), referValues);
    ret.push_back(colorCastTable);

    if (!colorCastTable.empty())
    {
        file.open(path + "cc" + std::to_string(dpi) + "b.dat", std::ios::binary | std::ios::out);
        file.write((char*)colorCastTable.data, 256);
        file.close();
    }

    return ret;
}

cv::Mat readZipCorrect(uchar* data, int len, int blackLevel)
{
    std::vector<std::vector<uchar>> colors;
    int index = 0, cols = len / 26;
    cv::Mat lut(cols, 256, CV_8UC1);
    for (size_t i = 0; i < cols; i++)
    {
        std::vector<uchar> color;
        for (size_t j = 0; j < 26; j++, index++)
            color.push_back(data[index]);
        colors.push_back(color);
    }

    for (int i = 0; i < colors.size(); i++)
        for (int j = 0; j < blackLevel; j++)
            colors[i].erase(colors[i].begin());

    double ss = 250.0 / (colors[0].size() - 1);

    for (int i = 0; i < cols; i++)
    {
        uchar* ptr = lut.data + 256 * i;
        memset(ptr, 255, 256);
        memset(ptr, 0, 100);

        for (int y = 1; y < colors[i].size(); y++)
        {
            int index_0 = colors[i][y - 1];
            int index_1 = colors[i][y];

            double value_0 = (y - 1) * ss;
            double value_1 = cv::min(y * ss, 255.0);
            double step = (double)(value_1 - value_0) / (double)(index_1 - index_0);
            for (int x = index_0; x < index_1; x++)
                ptr[x] = cv::min((int)(value_0 + (x - index_0) * step), 255);
        }
    }

    const int channels = (lut.rows % 408 == 0) ? 408 : 432;

    for (size_t i = 0, block = lut.rows / channels; i < block; i++)
    {
        cv::Mat lutROI = lut(cv::Rect(0, i * channels, 256, channels));
        cv::Mat tran;
        cv::transpose(lutROI, tran);
        memcpy(lutROI.data, tran.data, tran.total());
    }

    cv::Mat lutMat(lut.rows / channels, 256, CV_8UC(channels));
    memcpy(lutMat.data, lut.data, lut.total());
    return lutMat;
}