#include <iostream>
#include <opencv2/highgui.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/core/ocl.hpp>
#include <opencv2/features2d.hpp>
#include <opencv2/viz/types.hpp>
#include <math.h>
#include <unistd.h>
#include <vector>
#include <string>
#include <sstream>
#include <algorithm> 

#include "argpopt.h"

bool verbose = false;

void cd_to_exe_dir( char *argv[] )
{
	std::string path = argv[0];
	int ii = path.length();
	while ( !( path[ii] == '/' || path[ii] == '\\' ) && ii > 0 )
	{
		ii--;
	}
	path.erase( ii, 100 );
	chdir( path.c_str() );
}

std::vector<cv::Mat> loadImages(char** fileNames)
{ 
	std::vector<cv::Mat> images;
	for(size_t i = 0; fileNames[i]; ++i)
	{
		cv::Mat tmpImage = cv::imread(fileNames[i]);
		if(tmpImage.data)images.push_back(tmpImage);
		else std::cout<<"can not read image "<<i<<" from "<<fileNames[i]<<'\n';
	}
	return images;
}


std::vector< std::vector<cv::Point2f > > sortPointsIntoRows(std::vector<cv::Point2f>& points)
{
	
	if(points.size() < 6) return std::vector< std::vector<cv::Point2f> >();
	
	struct 
	{
		bool operator()(const cv::Point2f& a, const cv::Point2f& b) const
		{   
			return sqrt(a.y*a.y+a.x*a.x) < sqrt(b.y*b.y+b.x*b.x);
		}   
	} lessDist;
	
	std::vector<cv::Point2f>::iterator topLeftIt = std::min_element(points.begin(), points.end(), lessDist);
	std::vector<cv::Point2f>::iterator bottomRightIt = std::max_element(points.begin(), points.end(), lessDist);
	
	cv::Point2f topLeft(topLeftIt[0]);
	cv::Point2f bottomRight(bottomRightIt[0]);
	
	std::cout<<"topLeft "<<topLeft.x<<' '<<topLeft.y<<'\n';
	std::cout<<"bottomRight "<<bottomRight.x<<' '<<bottomRight.y<<'\n';
	
	float fuzz = (bottomRight.x-topLeft.x)*0.01f;
	
	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;
}

void drawRows(cv::Mat& image, const std::vector< std::vector<cv::Point2f > >& rows)
{
	for(size_t i = 0; i < rows.size(); ++i)
	{
		for(size_t y = 0; y < rows[i].size(); ++y)
		{
			cv::circle(image, rows[i][y], 5, cv::viz::Color(128 * (i%3), 128 * ((i+1)%3), 128 * ((i+2)%3)));
		}
	}
}

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);
	}
	return ellipses;
}

void drawEllipses(cv::Mat& image, const std::vector<cv::RotatedRect>& ellipses )
{
	 for(const auto& ellipse : ellipses)cv::ellipse(image, ellipse, cv::viz::Color(128,128,128));
}


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

DisplacmentMap calcDisplacementMap(const std::vector< std::vector<cv::Point2f > >& rows, 
											   const std::vector<cv::RotatedRect>& elipses)
{
	if(rows.size() < 2 || rows.size() != elipses.size()) {
		std::cerr<<__func__<<": rows < 2 or rows != elipses\n";
		return DisplacmentMap();
	}
	
	DisplacmentMap displacmentmap;
	
	displacmentmap.destination.assign(rows.size(), std::vector<cv::Point2f>());
	displacmentmap.source = rows;
	
	for(int i = 0; i < displacmentmap.destination.size(); ++i)
		displacmentmap.source[i].reserve(rows[i].size());
	
	
	std::cout<<__func__<<": "<<elipses[1].center.y-elipses[0].center.y<<'\n';
	
	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();
	
	std::cout<<__func__<<": meanYdist "<<meanYdist<<'\n';
	
	for(int 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();
		
		std::cout<<__func__<<": Proc row "<<rowCnt<<'\n';
		
		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;
			std::cout<<__func__<<": normDist "<<normDist<<" tau "<<tau<<" xDest "<<xDest<<'\n';
			displacmentmap.destination[rowCnt].push_back(cv::Point2f(xDest,yDest));
		}
	}
	return displacmentmap;
}

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());
	
	std::cout<<__func__<<": xMinDist "<<xMinDist<<'\n';
	
	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;
}


double distance(const cv::Point2f& a, const cv::Point2f& b)
{
	return sqrt((a.y-b.y)*(a.y-b.y)+(a.x-b.x)*(a.x-b.x));
}

bool findClosest(size_t& xIndex, size_t& yIndex, 
				 cv::Point2f point, const std::vector< std::vector<cv::Point2f> >& array, 
				 float xTolerance, float yTolerance)
{	
	size_t rowBelow = 0;
	while(rowBelow < array.size() && array[rowBelow][0].y < point.y) ++rowBelow;
	
	if(rowBelow == array.size()) rowBelow = array.size()-1;
	
	int rightAbove = -1;
	int rightBelow = 0;
	
	while(rightBelow < array[rowBelow].size() && array[rowBelow][rightBelow].x < point.x ) ++rightBelow;
	
	float distRB = distance(point, array[rowBelow][rightBelow]);
	float distLB = rightBelow > 0 ? distance(point, array[rowBelow][rightBelow-1]) : std::numeric_limits<float>::max();
	float distRA = std::numeric_limits<float>::max();
	float distLA = std::numeric_limits<float>::max();
	
	if(rowBelow > 0)
	{
		while(rightAbove < array[rowBelow].size() && array[rowBelow-1][rightAbove].x < point.x ) ++rightAbove;
		distRA = distance(point, array[rowBelow-1][rightAbove]);
		if(rightAbove > 0)  distLA = distance(point, array[rowBelow-1][rightAbove-1]);
	}
	
	float* min = &distRB;
	if(distLB < *min) min = &distLB;
	if(distRA < *min) min = &distRA;
	if(distLA < *min) min = &distLA;
	
	if(min == &distRB)
	{
		yIndex = rowBelow;
		xIndex = rightBelow;
	} 
	else if(min == &distLB)
	{
		yIndex = rowBelow;
		xIndex = rightBelow-1;
	} 
	else if(min == &distRA)
	{
		yIndex = rowBelow-1;
		xIndex = rightBelow;
	} 
	else if(min == &distLA)
	{
		yIndex = rowBelow-1;
		xIndex = rightBelow-1;
	}
	return abs(array[yIndex][xIndex].x - point.x) < xTolerance && abs(array[yIndex][xIndex].y - point.y) < yTolerance;
}


void interpolateMissing(cv::Mat& mat)
{
	for(size_t y = 0; y < mat.rows; y++)
	{
		float* col = mat.ptr<float>(y);
		for(int x = 0; x < mat.cols; ++x)
		{
			if(col[x] < 0) 
			{
				int closestA = -1;
				int closestB = -1;
				int dist = std::numeric_limits<int>::max();
				for(int i = 0; i < mat.cols; ++i)
				{
					if(col[i] >= 0 && abs(i-x) <= dist)
					{
						closestB = closestA;
						closestA = i;
						dist = abs(i-x);
					}
				}
				float slope = (col[closestB] - col[closestA])/(closestB-closestA);
				col[x] = col[closestA] - (closestA-x)*slope;
			}
		}
	}
}

void fillMissing(cv::Mat& mat)
{
	for(size_t y = 0; y < mat.rows; y++)
	{
		float* col = mat.ptr<float>(y);
		for(int x = 0; x < mat.cols; ++x)
		{
			if(col[x] < 0) 
			{
				if(y > 0 && mat.at<float>(x,y-1) >= 0)
					col[x] = mat.at<float>(x,y-1);
				else if(y < mat.rows && mat.at<float>(x,y+1) >= 0)
					col[x] = mat.at<float>(x,y+1);
				if((x+1 < mat.cols && col[x] > col[x+1]) || (x > 0 && col[x] < col[x-1]))
					col[x] = -1;
			}
		}
	}
}

void sanityCheckMap(cv::Mat& mat, const float min, const float max)
{
	for(size_t y = 0; y < mat.rows; y++)
	{
		float* col = mat.ptr<float>(y);
		for(int x = 0; x < mat.cols; ++x)
		{
			if(col[x] < min) 
				col[x] = min;
			else if(col[x] > max)
				col[x] = max;
		}
	}
}

//dst(x,y) =  src(map_x(x,y),map_y(x,y))
void generateRemapMaps(const DisplacmentMap& map, cv::Mat& xMat, cv::Mat& yMat, const cv::Size& size, float fuzz = 1.3f)
{
	if(map.destination.size() < 2) 
	{
		std::cerr<<__func__<<": at least 2 rows are needed"<<std::endl;
		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();
	std::cout<<__func__<<": xMeanDist "<<xMeanDist<<'\n';
	
	float xMin = std::numeric_limits<float>::max();
	float xMax = std::numeric_limits<float>::min();
	
	for(auto& row: map.destination )
	{
		if(xMin > row.front().x )
			xMin = row.front().x;
		if(xMax < row.back().x)
			xMax = row.back().x;
	}
	
	std::cout<<__func__<<": Grid: xMin "<<xMin<<'\n';
	std::cout<<__func__<<": Grid: grid xMax "<<xMax<<'\n';
	
	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);
	
	std::cout<<"Grid: "<<xGridSize<<'x'<<map.destination.size()<<'\n';
	
	for(size_t 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;
			size_t yIndex;
			bool found = findClosest(xIndex, yIndex, 
									 cv::Point2f((x+1)*xMeanDist, (y)*yMeanDist), 
									 map.destination, xMeanDist/2, yMeanDist/2);
			std::cout<<__func__<<": found: "<<found<<' '<<xIndex<<"x"<<yIndex<<'\n';
			colx[x] = found ? map.source[yIndex][xIndex].x : -1;
			coly[x] = found ? map.source[yIndex][xIndex].y : -1;
		}
	}
	fillMissing(xMat);
	interpolateMissing(xMat);
	interpolateMissing(yMat);
	std::cout<<__func__<<": xMat \n"<<xMat<<'\n';
	std::cout<<__func__<<": yMat \n"<<yMat<<'\n';
	resize(xMat, xMat, size);
	resize(yMat, yMat, size);
}

std::vector<cv::Point2f> detectPoints(cv::Mat& image)
{
		cv::Mat gray;
		cv::cvtColor(image, gray, cv::COLOR_BGR2GRAY);
		
		//detect corners
		cv::Mat corners;
		cv::cornerHarris(gray, corners, 5, 5, 0.01);
		cv::normalize(corners, corners, 0, 255, cv::NORM_MINMAX, CV_32FC1, cv::Mat());
		cv::convertScaleAbs( corners, corners );
		cv::threshold(corners, corners, 40, 255, cv::THRESH_BINARY);
		
		cv::waitKey(0);
		cv::imshow( "Viewer", corners );
		
		//get middle of corners
		cv::SimpleBlobDetector::Params blobParams;
		blobParams.filterByArea = true;
		blobParams.minArea = 4;
		blobParams.maxArea = 50;
		blobParams.filterByColor = false;
		blobParams.blobColor = 255;
		blobParams.filterByInertia = false;
		blobParams.filterByConvexity = false;
		cv::Ptr<cv::SimpleBlobDetector> blobDetector = cv::SimpleBlobDetector::create(blobParams);
		std::vector<cv::KeyPoint> keypoints;
		blobDetector->detect(corners, keypoints);
		
		std::vector<cv::Point2f> points;
		cv::KeyPoint::convert(keypoints, points);
		
		return points;
}

int main(int argc, char* argv[])
{
	cv::ocl::setUseOpenCL(false);
	std::cout<<"UVOS Optical lubricant thikness mapper "<<argp_program_version<<std::endl;
	cd_to_exe_dir(argv);
	
	Config config;
    argp_parse(&argp, argc, argv, 0, 0, &config);
	
	std::vector<cv::Mat> inImages = loadImages(config.inFileNames);
	
	if(inImages.empty()) 
	{
		std::cout<<"Input images must be provided!\n";
		return -1;
	}
	
	for(auto& image : inImages)
	{
		cv::namedWindow( "Viewer", cv::WINDOW_NORMAL );
		cv::resizeWindow("Viewer", 960, 500);
		cv::imshow( "Viewer", image );
		std::vector<cv::Point2f > points = detectPoints(image);
		
		std::vector< std::vector<cv::Point2f > > rows = sortPointsIntoRows(points);
		std::vector< cv::RotatedRect > ellipses = fitEllipses(rows);
		
		cv::Mat pointsMat = cv::Mat::zeros(image.size(), CV_8UC3);
		
		drawEllipses(pointsMat, ellipses);
		
		std::cout<<"rows: "<<rows.size()<<'\n';
		
		
		DisplacmentMap dispMap = calcDisplacementMap(rows, ellipses);
		
		drawRows(pointsMat, dispMap.source);
		
		cv::waitKey(0);
		cv::imshow( "Viewer", pointsMat );
		
		cv::Mat dispPointsDest = cv::Mat::zeros(cv::Size(cv::Size(1000,1000)), CV_8UC3);
		drawRows(dispPointsDest, dispMap.destination);
		
		cv::waitKey(0);
		cv::imshow( "Viewer", dispPointsDest );
		
		cv::Mat xMat;
		cv::Mat yMat;
		generateRemapMaps(dispMap, xMat, yMat, cv::Size(1000,1000));
		
		cv::Mat remaped;
		cv::remap(image, remaped,  xMat, yMat, cv::INTER_LINEAR);
		cv::waitKey(0);
		cv::imshow( "Viewer", remaped );
		cv::waitKey(0);
		cv::destroyWindow("Viewer");  
	}
	
	return 0;
}