/**
* Lubricant Detecter
* Copyright (C) 2021 Carl Klemm
*
* This program is free software; you can redistribute it and/or
* modify it under the terms of the GNU General Public License
* version 3 as published by the Free Software Foundation.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
* GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public License
* along with this program; if not, write to the
* Free Software Foundation, Inc., 51 Franklin Street, Fifth Floor,
* Boston, MA  02110-1301, USA.
*/

#include "matutils.h"
#include <opencv2/core/ocl.hpp>
#include <iostream>
#include <math.h>
#include <assert.h>
#include "uvosunwrap/log.h"

void sanityCheckMap(cv::Mat& mat, const float min, const float max, float minValue, float maxValue)
{
	for(int 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] = minValue;
			else if(col[x] > max)
				col[x] = maxValue;
		}
	}
}

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

std::vector<cv::Point2f>::iterator getTopLeft(std::vector<cv::Point2f>& points)
{
	return std::min_element(points.begin(), points.end(), [](const cv::Point2f& a, const cv::Point2f& b){
		return sqrt(a.y*a.y+a.x*a.x) < sqrt(b.y*b.y+b.x*b.x);
	});
}

std::vector<cv::Point2f>::iterator getBottomRight(std::vector<cv::Point2f>& points)
{
	return std::max_element(points.begin(), points.end(), [](const cv::Point2f& a, const cv::Point2f& b){
		return sqrt(a.y*a.y+a.x*a.x) < sqrt(b.y*b.y+b.x*b.x);
	});
}

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, 
				 const 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].empty() && array[rowBelow][0].y < point.y) ++rowBelow;
	
	if(rowBelow == array.size()) rowBelow = array.size()-1;
	
	size_t rightAbove = 0;
	size_t 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-1].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;
	}
	else return false;
	return abs(array[yIndex][xIndex].x - point.x) < xTolerance && abs(array[yIndex][xIndex].y - point.y) < yTolerance;
}

bool findClosest(size_t& index, const cv::Point2f point, const std::vector<cv::Point2f>& array, float xTolerance, float yTolerance)
{
	float minDistance = std::numeric_limits<float>::max();

	bool found;
	for(size_t i = 0; i < array.size(); ++i)
	{
		const auto& arrayPoint = array[i];
		if(abs(arrayPoint.x - point.x < xTolerance) && 
		   abs(arrayPoint.y - point.y < yTolerance) &&
		   distance(point, arrayPoint) < minDistance)
		{
			index = i;
			found = true;
		}
	}
	return found;
}

void interpolateMissingOnX(cv::Mat& mat)
{
	assert(mat.type() == CV_32FC1);
	
	for(int 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(i != closestA && col[i] >= 0 && abs(i-x) <= dist)
					{
						closestB = closestA;
						closestA = i;
						dist = abs(i-x);
					}
				}
				if(closestA < 0 || closestB < 0)
				{
					closestA = -1;
					closestB = -1;
					dist = std::numeric_limits<int>::max();
					for(int i = mat.cols-1; i >= 0; --i)
					{
						if(i != closestA && 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 interpolateMissingOnY(cv::Mat& mat)
{
	assert(mat.type() == CV_32FC1);
	
	for(int x = 0; x < mat.cols; ++x)
	{
		for(int y = 0; y < mat.rows; y++)
		{
			if(mat.at<float>(y,x) < 0) 
			{
				int closestA = -1;
				int closestB = -1;
				int dist = std::numeric_limits<int>::max();
				for(int i = 0; i < mat.rows; i++)
				{
					if(i != closestA && mat.at<float>(i,x) >= 0 && abs(i-y) <= dist)
					{
						closestB = closestA;
						closestA = i;
						dist = abs(i-y);
					}
				}
				if(closestA < 0 || closestB < 0)
				{
					closestA = -1;
					closestB = -1;
					dist = std::numeric_limits<int>::max();
					for(int i = mat.rows-1; i >= 0; --i)
					{
						if(i != closestA && mat.at<float>(i,x) >= 0 && abs(i-y) <= dist)
						{
							closestB = closestA;
							closestA = i;
							dist = abs(i-y);
						}
					}
				}
				float slope = (mat.at<float>(closestB,x) - mat.at<float>(closestA,x))/(closestB-closestA);
				mat.at<float>(y,x) = mat.at<float>(closestA,x) - (closestA-y)*slope;
			}
		}
	}
}

bool findDeadSpace(const cv::Mat& mat, cv::Rect& roi)
{
	assert(mat.type() == CV_32FC1);
	
	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 > 0 ? left : 0;
	roi.y=top > 0 ? top : 0;
	roi.width = right-roi.x;
	roi.height = bottom-roi.y;
	
	if(roi.width < 0 || roi.height < 0)
		return false;
	
	Log(Log::DEBUG)<<"Reduced Roi: "<<roi;
	return true;
}

void removeSparseCols(cv::Mat& mat, float reject, bool front)
{
	assert(mat.type() == CV_32FC1);

	int count = 0;
	for(int y = 0; y < mat.rows; ++y)
	{
		if((front && mat.at<float>(y,0) >= 0) || (!front && mat.at<float>(y,mat.cols-1) >= 0))
			++count;
	}
	cv::Rect roi;
	bool rej = (count < mat.rows*reject);
	roi.x=front ? rej : 0;
	roi.y=0;
	roi.width = mat.cols - rej;
	roi.height = mat.rows;
	mat = mat(roi);
	if(rej)
		removeSparseCols(mat, reject, front);
}

static float linInterpolate(float A, float B, float x)
{
	return A * (1.0f - x) + B * x;
}

template <typename Type> static Type resizePointBilinear(const cv::Mat& inImage, float scaledX, float scaledY, 
														 float xFrac, float yFrac)
{
	int u0 = static_cast<int>(scaledX);
	int v0 = static_cast<int>(scaledY);
	int u1 = std::min(inImage.cols-1, static_cast<int>(scaledX)+1);
	int v1 = v0;
	int u2 = u0;
	int v2 = std::min(inImage.rows-1, static_cast<int>(scaledY)+1);
	int u3 = std::min(inImage.cols-1, static_cast<int>(scaledX)+1);
	int v3 = std::min(inImage.rows-1, static_cast<int>(scaledY)+1);

	float col0 = linInterpolate(inImage.at<Type>(v0, u0), inImage.at<Type>(v1, u1), xFrac);
	float col1 = linInterpolate(inImage.at<Type>(v2, u2), inImage.at<Type>(v3, u3), xFrac);
	float value = linInterpolate(col0, col1, yFrac);

	return cv::saturate_cast<Type>(value);
}

template <typename Type> static void resize(const cv::Mat& inImage, cv::Mat& outImage) 
{
	float scaleY = (inImage.rows - 1) / static_cast<float>(outImage.rows - 1);
	float scaleX = (inImage.cols - 1) / static_cast<float>(outImage.cols - 1);

	for (int y = 0; y < outImage.rows; ++y) 
	{
		float scaledY = y * scaleY;
		float yFrac = scaledY - static_cast<int>(scaledY);

		for (int x = 0; x < outImage.cols; x++)
		{
			float scaledX = x * scaleX;
			float xFrac = scaledX - (static_cast<int>(scaledX));

			outImage.at<Type>(y, x) = resizePointBilinear<Type>(inImage, scaledX, scaledY, xFrac, yFrac);
		}
	}
}

bool bilinearResize(const cv::Mat& inImage, cv::Mat& outImage, const cv::Size size) 
{
	if (!(inImage.type() == CV_8U || inImage.type() == CV_32F) ||
	size.width < 2 || size.height < 2 || inImage.cols < 2 || inImage.rows < 2)
		return false;

	std::cout<<"size: "<<size.height<<' '<<size.width<<'\n';
	
	outImage = cv::Mat::zeros(size.height, size.width, inImage.type());

	if(inImage.type() == CV_8U)
		resize<uint8_t>(inImage, outImage);
	else if(inImage.type() == CV_32F)
		resize<float>(inImage, outImage);
	return true;
}