#include "unwrap.h"

#include <opencv2/highgui.hpp>
#include <opencv2/imgproc.hpp>
#include <math.h>
#include <vector>
#include <algorithm>
#include <iostream>

#include "matutils.h"
#include "drawing.h"
#include "log.h"

struct DisplacmentMap
{
	std::vector< std::vector<cv::Point2f> > destination;
	std::vector< std::vector<cv::Point2f> > source;
};

static std::vector< std::vector<cv::Point2f > > sortPointsIntoRows(std::vector<cv::Point2f>& points)
{
	
	if(points.size() < 6) return std::vector< std::vector<cv::Point2f> >();
	
	cv::Point2f topLeft(*getTopLeft(points));
	cv::Point2f bottomRight(*getBottomRight(points));
	
	Log(Log::DEBUG)<<"topLeft "<<topLeft.x<<' '<<topLeft.y;
	Log(Log::DEBUG)<<"bottomRight "<<bottomRight.x<<' '<<bottomRight.y;
	
	float fuzz = (bottomRight.x-topLeft.x)*0.05f;
	
	for(size_t i = 0; i < points.size(); ++i)
	{
		if(points[i].x < topLeft.x-fuzz || points[i].y < topLeft.y-fuzz || 
			points[i].x > bottomRight.x+fuzz || points[i].y > bottomRight.y+fuzz)
			points.erase(points.begin()+i);
	}
	
	std::sort(points.begin(), points.end(), [](const cv::Point2f& a, const cv::Point2f& b){return a.y < b.y;});
	
	double accumulator = points[0].y+points[1].y;
	size_t accuCount = 2;
	float sigDist = (bottomRight.y-topLeft.y) / 20;
	
	std::vector< std::vector<cv::Point2f> > result(1);
	result.back().push_back(points[0]);
	result.back().push_back(points[1]);
	
	for(size_t i = 2; i < points.size(); ++i)
	{
		if( points[i].y - accumulator/accuCount > sigDist )
		{
			if(points.size() - i - 3 < 0) break;
			else 
			{
				accumulator = points[i+1].y+points[i+2].y;
				accuCount = 2;
			}
			result.push_back(std::vector<cv::Point2f>());
		}
		result.back().push_back(points[i]);
		accumulator += points[i].y;
		++accuCount;
	}
	
	for(auto& row : result)  std::sort(row.begin(), row.end(), [](const cv::Point2f& a, const cv::Point2f& b){return a.x < b.x;});
	
	return result;
}


static float detimineXPitch(const std::vector<cv::Point2f>& row, float fuzz = 1.3f)
{
	std::vector<float> xRowDists;
	for(size_t i = 0; i < row.size()-1; ++i)
		xRowDists.push_back(abs(row[i+1].x-row[i].x));
	float xMinDist = *std::min(xRowDists.begin(), xRowDists.end());
	
	Log(Log::DEBUG)<<__func__<<": xMinDist "<<xMinDist;
	
	float meanXDistAccum = 0;
	size_t validCount = 0;
	for(size_t i = 0; i < xRowDists.size(); ++i)
	{
		if(xRowDists[i] < xMinDist*fuzz)
		{
			++validCount;
			meanXDistAccum+=xRowDists[i];
		}
	}
	return meanXDistAccum/validCount;
}


static std::vector<cv::RotatedRect> fitEllipses(std::vector< std::vector<cv::Point2f > >& rows, bool remove = true)
{
	if(rows.empty()) return std::vector<cv::RotatedRect>();
	
	std::vector<cv::RotatedRect> ellipses;
	ellipses.reserve(rows.size());
	for(size_t i = 0; i < rows.size(); ++i) 
	{
		if(rows[i].size() > 4) ellipses.push_back(cv::fitEllipse(rows[i]));
		else 
		{ 
			rows.erase(rows.begin()+i);
			i--;
		}
	}
	return ellipses;
}

static void thompsonTauTest(const std::vector<float>& in, std::vector<size_t>& outliers, float criticalValue)
{
	float mean = 0;
	size_t n = 0;
	for(size_t i = 0; i < in.size(); ++i)
	{
		bool found = false;
		for(size_t j = 0; j < outliers.size() && !found; ++j) 
			if(outliers[j]==i) found = true;
		if(!found)
		{
			mean+=in[i];
			++n;
		}
	}
	mean/=n;
	
	float sd = 0;
	for(size_t i = 0; i < in.size(); ++i)
	{
		bool found = false;
		for(size_t j = 0; j < outliers.size() && !found; ++j) 
			if(outliers[j]==i) found = true;
		if(!found) sd+=pow(in[i]-mean,2);
	}
	sd = sqrt(sd/(n-1));
	float rej = (criticalValue*(n-1))/(sqrt(n)*sqrt(n-2+criticalValue));
	bool removed = false;
	for(size_t i = 0; i < in.size() && !removed; ++i)
	{
		bool found = false;
		for(size_t j = 0; j < outliers.size() && !found; ++j) 
			if(outliers[j]==i) found = true;
		if(!found)
		{
			if(abs((in[i]-mean)/sd) > rej)
			{
				Log(Log::DEBUG)<<__func__<<": "<<i<<" is outlier mean: "<<mean<<" sd: "<<sd<<" n: "<<n<<'\n';
				outliers.push_back(i);
				removed = true;
			}
		}
	}
	if(removed) thompsonTauTest(in, outliers, criticalValue);
}

static bool sanityCheckElipses(std::vector< std::vector<cv::Point2f > >& rows, std::vector<cv::RotatedRect>& elipses)
{
	if(rows.size() != elipses.size() && elipses.size() > 1) return false;

	for(size_t i = 0; i < elipses.size(); ++i)
	{
		float angDiff = fmod(elipses[i].angle,90);
		if(angDiff < 90-5 && angDiff > 5)
		{
			elipses.erase(elipses.begin()+i);
			rows.erase(rows.begin()+i);
			--i;
		}
	}
	
	std::vector<float> widths;
	std::vector<float> heights;
	
	for(auto& elipse : elipses)
	{
		widths.push_back(elipse.size.width);
		heights.push_back(elipse.size.height);
	}
	
	std::vector<size_t> outliersW;
	std::vector<size_t> outliersH;
	thompsonTauTest(widths, outliersW, 2);
	thompsonTauTest(heights, outliersH, 2);
	
	std::vector< std::vector<cv::Point2f > > rowsReduced;
	std::vector<cv::RotatedRect> elipsesReduced;

	for(size_t i = 0; i < elipses.size(); ++i)
	{
		bool found = false;
		for(size_t j = 0; j < outliersW.size() && !found; ++j) 
			if(outliersW[j]==i) found = true;
		for(size_t j = 0; j < outliersH.size() && !found; ++j) 
			if(outliersH[j]==i) found = true;
		if(!found)
		{
			rowsReduced.push_back(rows[i]);
			elipsesReduced.push_back(elipses[i]);
		}
	}
	elipses = elipsesReduced;
	rows = rowsReduced;
	return true;
}

static DisplacmentMap calcDisplacementMap(const std::vector< std::vector<cv::Point2f > >& rows, 
											   const std::vector<cv::RotatedRect>& elipses)
{
	if(rows.size() < 2 || rows.size() != elipses.size()) {
		Log(Log::ERROR)<<__func__<<": rows < 2 or rows != elipses";
		return DisplacmentMap();
	}
	
	DisplacmentMap displacmentmap;
	
	displacmentmap.destination.assign(rows.size(), std::vector<cv::Point2f>());
	displacmentmap.source = rows;
	
	for(size_t i = 0; i < displacmentmap.destination.size(); ++i)
	{
		displacmentmap.destination[i].reserve(rows[i].size());
	}
	
	float meanYdist = 0;
	size_t j = 0;
	while(j < elipses.size()-1)
	{
		meanYdist += elipses[j+1].center.y-elipses[j].center.y;
		++j;
	}
	meanYdist = meanYdist/elipses.size();
	
	Log(Log::DEBUG)<<__func__<<": meanYdist "<<meanYdist;
	
	for(size_t rowCnt = 0; rowCnt < rows.size(); ++rowCnt)
	{
		const std::vector<cv::Point2f>& row = rows[rowCnt];
		const cv::RotatedRect& elipse = elipses[rowCnt];
				
		cv::Rect_<float> boundingRect = elipse.boundingRect2f();
		
		Log(Log::DEBUG)<<__func__<<": Proc row "<<rowCnt;
		
		for(size_t i = 0; i < row.size(); ++i)
		{
			float yDest = rowCnt*meanYdist;
			
			double normDist = ((row[i].x - boundingRect.x)/boundingRect.width-0.5)*2;
			double tau = asin(normDist);
			float xDest = (((2*tau)/M_PI)*500)+500;
			Log(Log::DEBUG)<<__func__<<": normDist "<<normDist<<" tau "<<tau<<" xDest "<<xDest;
			displacmentmap.destination[rowCnt].push_back(cv::Point2f(xDest,yDest));
		}
	}
	return displacmentmap;
}

static void removeSparseCollums(cv::Mat& mat)
{
	int front = 0;
	int back = 0;
	for(int y = 0; y < mat.rows; ++y)
	{
		if(mat.at<float>(y,0) >= 0) ++front;
		if(mat.at<float>(y,mat.cols-1) >= 0) ++back;
	}
	Log(Log::DEBUG)<<__func__<<" front: "<<front<<" back "<<back;
	cv::Rect roi;
	bool frontRej = (front < mat.rows/2);
	bool backRej  = (back < mat.rows/2);
	roi.x=frontRej;
	roi.y=0;
	roi.width = mat.cols - backRej - frontRej;
	roi.height = mat.rows;
	mat = mat(roi);
	if(frontRej || backRej) removeSparseCollums(mat);
}


static void generateRemapMaps(const DisplacmentMap& map, cv::Mat& xMat, cv::Mat& yMat, float fuzz = 1.3f)
{
	if(map.destination.size() < 2) 
	{
		Log(Log::ERROR)<<__func__<<": at least 2 rows are needed";
		return;
	}
	
	float yMeanDist = map.destination[1][0].y-map.destination[0][0].y;
	float xMeanDist = 0;
	for(size_t i = 0; i < map.destination.size(); ++i)
	{
		xMeanDist+=detimineXPitch(map.destination[i]);
	}
	xMeanDist/=map.destination.size();
	Log(Log::DEBUG)<<__func__<<": xMeanDist "<<xMeanDist;
	
	float xMin = std::numeric_limits<float>::max();
	float xMeanMin = 0;
	float xMax = std::numeric_limits<float>::min();
	
	for(auto& row: map.destination )
	{
		xMeanMin+=row.front().x;
		if(xMin > row.front().x )
			xMin = row.front().x;
		if(xMax < row.back().x)
			xMax = row.back().x;
	}
	xMeanMin = xMeanMin / map.destination.size();
	
	Log(Log::DEBUG)<<__func__<<": Grid: xMin "<<xMin;
	Log(Log::DEBUG)<<__func__<<": Grid: grid xMax "<<xMax;
	Log(Log::DEBUG)<<__func__<<": Grid: grid xMeanMin "<<xMeanMin;
	
	size_t xGridSize = static_cast<size_t>(std::lround(abs((xMax-xMin)/xMeanDist))+1);
	xMat = cv::Mat::zeros(cv::Size(xGridSize, map.destination.size()), CV_32FC1);
	yMat = cv::Mat::zeros(cv::Size(xGridSize, map.destination.size()), CV_32FC1);
	
	Log(Log::DEBUG)<<"Grid: "<<xGridSize<<'x'<<map.destination.size();
	
	for(int y = 0; y < xMat.rows; y++)
	{
		float* colx = xMat.ptr<float>(y);
		float* coly = yMat.ptr<float>(y);
		for(int x = 0; x < xMat.cols; x++)
		{
			size_t xIndex = 0;
			size_t yIndex = 0;
			bool found = findClosest(xIndex, yIndex, 
									 cv::Point2f((x)*xMeanDist+xMin, (y)*yMeanDist), 
									 map.destination, (2*xMeanDist)/5, (2*yMeanDist)/5);
			Log(Log::DEBUG)<<__func__<<": found: "<<found<<' '<<xIndex<<'x'<<yIndex<<" at: "<<(x)*xMeanDist+xMin<<'x'<<(y)*yMeanDist;
			colx[x] = found ? map.source[yIndex][xIndex].x : -1;
			coly[x] = found ? map.source[yIndex][xIndex].y : -1;
		}
	}
	Log(Log::DEBUG)<<__func__<<": xMat raw\n"<<xMat;
	removeSparseCollums(xMat);
	removeSparseCollums(yMat);
	Log(Log::DEBUG)<<__func__<<": xMat rejcted\n"<<xMat;
	fillMissing(xMat);
	Log(Log::DEBUG)<<__func__<<": xMat filled\n"<<xMat;
	interpolateMissing(xMat);
	Log(Log::DEBUG)<<__func__<<": yMat raw \n"<<yMat;
	interpolateMissing(yMat);
	Log(Log::INFO)<<__func__<<": xMat \n"<<xMat;
	Log(Log::INFO)<<__func__<<": yMat \n"<<yMat;
}

static void generateRemapMaps(const std::vector<cv::Point2f>& points, const std::vector<cv::Point2i>& coordinates, cv::Mat& xMat, cv::Mat& yMat)
{
	if(points.size() < 6 || coordinates.size() != points.size()) 
	{
		Log(Log::ERROR)<<__func__<<": at least 6 points are needed";
		return;
	}
	
	int xMin = std::numeric_limits<int>::max();
	int xMax = std::numeric_limits<int>::min();
	int yMin = std::numeric_limits<int>::max();
	int yMax = std::numeric_limits<int>::min();
	
	for(auto& point : coordinates)
	{
		Log(Log::DEBUG)<<"point: "<<point.x<<'x'<<point.y;
		
		if(point.x > xMax) xMax = point.x;
		else if(point.x < xMin) xMin = point.x;
		if(point.y > yMax) yMax = point.y;
		else if(point.y < yMin) yMin = point.y;
	}
	
	size_t xGridSize = xMax-xMin+1;
	size_t yGridSize = yMax-yMin+1;
	xMat = cv::Mat::zeros(cv::Size(xGridSize, yGridSize), CV_32FC1);
	yMat = cv::Mat::zeros(cv::Size(xGridSize, yGridSize), CV_32FC1);
	
	Log(Log::DEBUG)<<"Grid: "<<xGridSize<<'x'<<yGridSize;
	
	for(int y = 0; y < xMat.rows; y++)
	{
		float* colx = xMat.ptr<float>(y);
		float* coly = yMat.ptr<float>(y);
		for(int x = 0; x < xMat.cols; x++)
		{
			colx[x] = -1;
			coly[x] = -1;
			for(size_t i = 0; i < coordinates.size(); ++i)
			{
				if(coordinates[i] == cv::Point2i(x+xMin,y+yMin))
				{
					colx[x] = points[i].x;
					coly[x] = points[i].y;
					break;
				}
			}
				
			
		}
	}
	Log(Log::DEBUG)<<__func__<<": xMat raw\n"<<xMat;
	removeSparseCollums(xMat);
	removeSparseCollums(yMat);
	Log(Log::DEBUG)<<__func__<<": xMat rejcted\n"<<xMat;
	fillMissing(xMat);
	Log(Log::DEBUG)<<__func__<<": xMat filled\n"<<xMat;
	interpolateMissing(xMat);
	Log(Log::DEBUG)<<__func__<<": yMat raw \n"<<yMat;
	interpolateMissing(yMat);
	Log(Log::INFO)<<__func__<<": xMat \n"<<xMat;
	Log(Log::INFO)<<__func__<<": yMat \n"<<yMat;
}

bool createRemapMap(cv::Mat& image, std::vector<cv::Point2f> points, const std::vector<cv::Point2i>& coordinates, 
					const std::string& fileName, bool verbose)
{

	cv::Mat xMat;
	cv::Mat yMat;
	
	if(coordinates.empty())
	{
		std::vector< std::vector<cv::Point2f > > rows = sortPointsIntoRows(points);
		Log(Log::INFO)<<"Found "<<rows.size()<<" rows";
		if(rows.size() < 2)
		{
			Log(Log::ERROR)<<"Error creating map, insufficant rows detected";
			return false;
		}
		
		if(verbose)
		{
			cv::Mat pointsMat = cv::Mat::zeros(image.size(), CV_8UC3);
			drawRows(pointsMat, rows);
			cv::imshow( "Viewer", pointsMat );
			cv::waitKey(0);
		}

		std::vector< cv::RotatedRect > ellipses = fitEllipses(rows);
		Log(Log::INFO)<<"Found "<<ellipses.size()<<" ellipses. rows reduced to "<<rows.size();
		if(ellipses.size() < 3)
		{
			Log(Log::ERROR)<<"Error creating map, insufficant ellipses detected";
			return false;
		}
		
		if(verbose)
		{
			cv::Mat pointsMat = cv::Mat::zeros(image.size(), CV_8UC3);
			drawRows(pointsMat, rows);
			drawEllipses(pointsMat, ellipses);
			cv::imshow( "Viewer", pointsMat );
			cv::waitKey(0);
		}
		
		sanityCheckElipses(rows, ellipses);
		
		if(verbose)
		{
			cv::Mat pointsMat = cv::Mat::zeros(image.size(), CV_8UC3);
			drawRows(pointsMat, rows);
			drawEllipses(pointsMat, ellipses);
			cv::imshow( "Viewer", pointsMat );
			cv::waitKey(0);
		}

		DisplacmentMap dispMap = calcDisplacementMap(rows, ellipses);
		if(dispMap.destination.size() < 2)
		{
			Log(Log::ERROR)<<"Error creating map, unable to calculate destination points";
			return false;
		}
		
		cv::Mat dispPointsDest = cv::Mat::zeros(cv::Size(cv::Size(1000,1000)), CV_8UC3);
		drawRows(dispPointsDest, dispMap.destination);
		
		if(verbose)
		{
			cv::imshow( "Viewer", dispPointsDest );
			cv::waitKey(0);
		}
		generateRemapMaps(dispMap, xMat, yMat);
	}
	else
	{
		generateRemapMaps(points, coordinates, xMat, yMat);
	}
	
	sanityCheckMap(xMat, 0, image.cols-1, -1, -1);
	sanityCheckMap(yMat, 0, image.rows-1, -1, -1);
	fillMissing(xMat);
	interpolateMissing(xMat);
	interpolateMissing(yMat);
	sanityCheckMap(xMat, 0, image.cols-1, 0, image.cols-1);
	sanityCheckMap(yMat, 0, image.rows-1, 0, image.rows-1);
	
	if(xMat.cols < 3 || xMat.rows < 3)
	{
		Log(Log::ERROR)<<"Error creating map, to few points with high confidence";
		return false;
	}
	
	cv::FileStorage matf(fileName+".mat", cv::FileStorage::WRITE );
	matf<<"xmat"<<xMat<<"ymat"<<yMat;
	matf.release();
	
	Log(Log::INFO)<<"Unwrap maps saved to "<<fileName<<".mat";
	
	if(verbose)
	{
		cv::Mat remaped;
		applyRemap(image, remaped, xMat, yMat, 800);
		cv::imshow( "Viewer", remaped );
		cv::waitKey(0);
	}
	return true;
}

bool findDeadSpace(const cv::Mat& mat, cv::Rect& roi)
{
	if(mat.cols < 2 || mat.rows < 2) return false;
	
	float centerTop = mat.at<float>(0,mat.cols/2);
	float centerLeft = mat.at<float>(mat.rows/2,0);
	float centerBottom = mat.at<float>(mat.rows-1,mat.cols/2);
	float centerRight = mat.at<float>(mat.rows/2,mat.cols-1);
	
	int left = 0;
	int right = mat.cols-1;
	int top = 0;
	int bottom = mat.rows-1;
	for(; left < mat.cols && mat.at<float>(mat.rows/2,left) == centerLeft; ++left);
	for(; top < mat.rows && mat.at<float>(top,mat.cols/2) == centerTop; ++top);
	for(; right > left && mat.at<float>(mat.rows/2,right) == centerRight; --right);
	for(; bottom > top && mat.at<float>(bottom,mat.cols/2) == centerBottom; --bottom);
	
	roi.x=left;
	roi.y=top;
	roi.width = right-left;
	roi.height = bottom-top;
	return true;
}

void applyRemap(cv::Mat& image, cv::Mat& out, const cv::Mat& xmap, const cv::Mat& ymap, unsigned int outputSize)
{
	cv::Mat xMapResized;
	cv::Mat yMapResized;
	cv::resize(xmap, xMapResized, cv::Size(outputSize,outputSize*(xmap.rows/(double)xmap.cols)), cv::INTER_CUBIC);
	cv::resize(ymap, yMapResized, cv::Size(outputSize,outputSize*(xmap.rows/(double)xmap.cols)), cv::INTER_CUBIC);
	cv::Rect roi;
	cv::Mat xMapRed;
	cv::Mat yMapRed;
	if(!findDeadSpace(xMapResized, roi))
	{
		xMapRed = xMapResized(roi);
		yMapRed = yMapResized(roi);
	}
	else
	{
		xMapRed = xMapResized;
		yMapRed = yMapResized;
	}
	cv::remap(image, out,  xMapRed, yMapRed, cv::INTER_LINEAR);
}