/**
* 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 "uvosunwrap/unwrap.h"

#include <opencv2/highgui.hpp>
#include <opencv2/stitching/detail/seam_finders.hpp>
#include <opencv2/stitching/detail/blenders.hpp>
#include <opencv2/imgproc.hpp>
#include <math.h>
#include <vector>
#include <algorithm>
#include <iostream>

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

static void sortIntoRemapMaps(const std::vector<DetectedPoint>& points, cv::Mat& xMat, cv::Mat& yMat, cv::Point2i& topLeftCoordinate)
{
	if(points.size() < 6) 
	{
		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 : points)
	{
		Log(Log::DEBUG)<<"point: "<<point.coordinate.x<<'x'<<point.coordinate.y;
		
		if(point.coordinate.x > xMax) xMax = point.coordinate.x;
		else if(point.coordinate.x < xMin) xMin = point.coordinate.x;
		if(point.coordinate.y > yMax) yMax = point.coordinate.y;
		else if(point.coordinate.y < yMin) yMin = point.coordinate.y;
	}
	
	topLeftCoordinate.x = xMin;
	topLeftCoordinate.y = yMin;
	
	size_t xGridSize = xMax-xMin+1;
	size_t yGridSize = yMax-yMin+1;
	xMat.create(cv::Size(xGridSize, yGridSize), CV_32FC1);
	yMat.create(cv::Size(xGridSize, yGridSize), CV_32FC1);
	xMat = -1;
	yMat = -1;
	
	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++)
		{
			for(size_t i = 0; i < points.size(); ++i)
			{
				if(points[i].coordinate == cv::Point2i(x+xMin,y+yMin))
				{
					colx[x] = points[i].point.x;
					coly[x] = points[i].point.y;
					break;
				}
			}
				
			
		}
	}

}

bool createRemapMap(const cv::Mat& image, RemapMap& out, const std::vector<DetectedPoint>& points, bool verbose)
{	
	sortIntoRemapMaps(points, out.xMat, out.yMat, out.topLeftCoordinate);
	
	Log(Log::DEBUG)<<__func__<<": xMat raw\n"<<out.xMat;
	removeEmptyFrontBackCols(out.xMat);
	removeEmptyFrontBackCols(out.yMat);
	deleteEmptyCols(out.xMat);
	deleteEmptyCols(out.yMat);
	Log(Log::DEBUG)<<__func__<<": xMat rejcted\n"<<out.xMat;
	fillMissing(out.xMat);
	Log(Log::DEBUG)<<__func__<<": xMat filled\n"<<out.xMat;
	interpolateMissing(out.xMat);
	interpolateMissing(out.yMat);

	sanityCheckMap(out.xMat, 0, image.cols-1, -1, -1);
	sanityCheckMap(out.yMat, 0, image.rows-1, -1, -1);
	fillMissing(out.xMat);
	interpolateMissing(out.xMat);
	interpolateMissing(out.yMat);
	sanityCheckMap(out.xMat, 0, image.cols-1, 0, image.cols-1);
	sanityCheckMap(out.yMat, 0, image.rows-1, 0, image.rows-1);
	
	Log(Log::INFO)<<__func__<<": xMat \n"<<out.xMat;
	Log(Log::INFO)<<__func__<<": yMat \n"<<out.yMat;
	
	if(out.xMat.cols < 3 || out.xMat.rows < 3)
	{
		Log(Log::ERROR)<<"Error creating map, to few points with high confidence";
		return false;
	}
	
	if(verbose)
	{
		RemapedImage remaped = applyRemap(image, out);
		cv::imshow( "Viewer", remaped.image );
		cv::waitKey(0);
	}

	return true;
}

bool saveRemapMap(const RemapMap& map, const std::string& fileName)
{
	cv::FileStorage matf(fileName, cv::FileStorage::WRITE );
	matf<<"xmat"<<map.xMat<<"ymat"<<map.yMat<<"origin"<<map.topLeftCoordinate;
	matf.release();
	Log(Log::INFO)<<"Unwrap maps saved to "<<fileName;
	return true;
}

RemapMap loadRemapMap(const std::string& fileName)
{
	cv::FileStorage fs(fileName, cv::FileStorage::READ);
	if (!fs.isOpened())
	{
		Log(Log::ERROR)<<"could not open maps file "<<fileName;
		return RemapMap();
	}
	RemapMap map;
	fs["xmat"]>>map.xMat;
	fs["ymat"]>>map.yMat;
	fs["origin"]>>map.topLeftCoordinate;
	return map;
}

RemapedImage applyRemap(const cv::Mat& image, const RemapMap &map)
{
	RemapedImage out;
	cv::Mat xMapResized;
	cv::Mat yMapResized;
	const cv::Size outputSize(map.outputCellSize*map.xMat.cols,map.outputCellSize*map.xMat.rows);
	
	cv::resize(map.xMat, xMapResized, outputSize, cv::INTER_LINEAR);
	cv::resize(map.yMat, yMapResized, outputSize, cv::INTER_LINEAR);
	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.image,  xMapRed, yMapRed, cv::INTER_LINEAR);
	out.origin = cv::Point2i(map.topLeftCoordinate.x*map.outputCellSize, map.topLeftCoordinate.y*map.outputCellSize);
	return out;
}

cv::Mat simpleStich(const std::vector<RemapedImage>& images)
{
	if(images.size() < 1)
	return cv::Mat();
	
	cv::Size outputSize(0,0);
	for(auto& image : images)
	{
		if(outputSize.width < image.image.cols+image.origin.x)
			outputSize.width = image.image.cols+image.origin.x;
		if(outputSize.height < image.image.rows+image.origin.y)
			outputSize.height = image.image.rows+image.origin.y;
		
		Log(Log::DEBUG)<<"image: "<<image.image.rows<<'x'<<image.image.cols<<" at "<<image.origin.x<<'x'<<image.origin.y;
	}
	
	Log(Log::DEBUG)<<"outputSize: "<<outputSize;
	
	cv::Mat out(outputSize, images[0].image.type(), cv::Scalar::all(0));
	
	for(auto& image : images)
	{
		cv::Rect roi(image.origin, image.image.size());
		image.image.copyTo(out(roi));
	}
	
	return out;
}

cv::Mat stich(const std::vector<RemapedImage>& images, bool seamAdjust)
{
	std::vector<cv::Mat> masks(images.size());
	std::vector<cv::Point> corners(images.size());
	std::vector<cv::Size> sizes(images.size());
	for (size_t i = 0; i < images.size(); i++)
	{
		masks[i].create(images[i].image.size(), CV_8U);
		masks[i].setTo(cv::Scalar::all(255));
		corners[i] = images[i].origin;
		sizes[i] = images[i].image.size();
	}
	
	
	if(seamAdjust)
	{
		std::vector<cv::UMat> images32f(images.size());
		std::vector<cv::UMat> masksUmat(images.size());
		for (size_t i = 0; i < images.size(); i++) 
		{
			images[i].image.convertTo(images32f[i], CV_32F);
			masks[i].copyTo(masksUmat[i]);
		}
		cv::detail::VoronoiSeamFinder seamFinder;
		seamFinder.find(images32f, corners, masksUmat);
		for (size_t i = 0; i < images.size(); i++) 
		{
			masksUmat[i].copyTo(masks[i]);
			masksUmat[i].release();
			images32f[i].release();
		}
		images32f.clear();
		masksUmat.clear();
	}

	cv::Ptr<cv::detail::Blender> blender;
	blender = cv::detail::Blender::createDefault(cv::detail::Blender::Blender::MULTI_BAND, false);
	cv::detail::MultiBandBlender* mb = dynamic_cast<cv::detail::MultiBandBlender*>(blender.get());
	mb->setNumBands(5);
	blender->prepare(corners, sizes);
	for (size_t i = 0; i < images.size(); i++)
	{
		blender->feed(images[i].image, masks[i], corners[i]);
	}
	cv::Mat result;
	cv::Mat result_mask;
	blender->blend(result, result_mask);
	
	return result;
}