#include <opencv2/core.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <string>
#include <iostream>
#include <array>
#include <fstream>

using namespace cv;
using namespace std;

array<Mat, 3> getHistograms(const Mat &img);
array<Mat, 3> getCumulativeHistograms(const Mat &img);
array<array<uint16_t, 256>, 3> getLUT(const Mat &img, const array<Mat, 3> &refCumHist);
Mat doLUT(const string &filename, const array<Mat, 3> &refCumHist);

int main()
{
	const constexpr size_t digitCount = 5;
	const constexpr size_t frameCount = 20505;
	auto refHistos = getCumulativeHistograms(imread("Frame00001.png", IMREAD_COLOR));

	for (size_t i = 2; i < frameCount; i++)
	{
		clog << "Processing frame: " << i << endl;
		string filename{ "Frame" };
		string numeric{ to_string(i) };

		filename += string(digitCount - numeric.size(), '0') + numeric + ".png";

		const auto lutted = doLUT(filename, refHistos);
		imwrite("processed/" + filename, lutted);
	}

	return 0;
}

array<Mat, 3> getHistograms(const Mat &img)
{
	array<Mat, 3> compPlanes;
	split(img, compPlanes);

	const int histSize = 256;

	const float range[]{ 0.f, 256.f };
	const float* histRange{ range };
	bool uniform = true, accumulate = false;

	array<Mat, 3> histos;

	for (size_t i = 0u; i < 3; i++)
		calcHist(&compPlanes[i], 1, nullptr, Mat{}, histos[i], 1, &histSize, &histRange, uniform, accumulate);

	for (size_t histI = 0u; histI < histos.size(); histI++)
	{
		double histTotal = 0.0;

		for (size_t i = 0u; i < histos[histI].rows; i++)
			histTotal += histos[histI].at<float>(i);

		for (size_t i = 0u; i < histos[histI].rows; i++)
			histos[histI].at<float>(i) /= histTotal;
	}

	return histos;
}

array<Mat, 3> getCumulativeHistograms(const Mat &img)
{
	auto histos = getHistograms(img);

	for (size_t histI = 0; histI < histos.size(); histI++)
	{
		double cumulativeTotal = 0.0;

		for (size_t i = 0; i < histos[histI].rows; i++)
		{
			cumulativeTotal += histos[histI].at<float>(i);
			histos[histI].at<float>(i) = cumulativeTotal;
		}
	}

	return histos;
}

array<array<uint16_t, 256>, 3> getLUT(const Mat &img, const array<Mat, 3> &refCumHist)
{
	const auto cumulativeHistos = getCumulativeHistograms(img);
	array<array<uint16_t, 256>, 3> luts{};

	auto comp = [](float a, float b){ return (b - a) > 1.0e-6f; };

	for (size_t lut = 0; lut < luts.size(); lut++)
	{
		uint16_t targetBin = 1;

		for (uint16_t bin = 0; bin < luts[lut].size(); bin++)
		{
			while (comp(refCumHist[lut].at<float>(targetBin), cumulativeHistos[lut].at<float>(bin)) && targetBin < luts[lut].size())
				targetBin++;

			luts[lut][bin] = targetBin - 1;
		}
	}

	return luts;
}

Mat doLUT(const string &filename, const array<Mat, 3> &refCumHist)
{
	Mat img{ imread(filename, IMREAD_COLOR) };

	if (!img.data)
	{
		cerr << "Failed to load image \'" << filename << "\'" << endl;
		return {};
	}

	const auto luts = getLUT(img, refCumHist);

	for (size_t x = 0; x < img.rows; x++)
	{
		for (size_t y = 0; y < img.cols; y++)
		{
			for (size_t c = 0; c < 3; c++)
				img.at<Vec3b>(x, y)[c] = luts[c][img.at<Vec3b>(x, y)[c]];
		}
	}

	return img;
}