/**
* 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 <iostream>
#include <opencv2/highgui.hpp>
#include <opencv2/core/ocl.hpp>
#include <fcntl.h>
#include <unistd.h>
#include <math.h> 
#include <vector>
#include <string> 
#include "argpopt.h"
#include "uvosunwrap/unwrap.h"
#include "uvosunwrap/bgremoval.h"
#include "uvosunwrap/normalize.h"
#include "uvosunwrap/log.h"
#include "uvosunwrap/charuco.h"
#include "uvosunwrap/harris.h"
#include "uvosunwrap/curve.h"

#define IMREAD_SIZE pow(2, 20)

enum {
	CREATE_CHARUCO,
	CREATE_MAP,
	APPLY_MAP,
	APPLY_CURVE,
	CREATE_CURVE,
	SHOW_IMAGE,
	EXIT
};

int selectOperation(char** opt)
{
	if(!opt || !opt[0])
		return -1;
	else if(strcmp(opt[0], "create" ) == 0)
		return CREATE_MAP;
	else if(strcmp(opt[0], "apply" ) == 0)
		return APPLY_MAP;
	else if(strcmp(opt[0], "curve" ) == 0)
		return APPLY_CURVE;
	else if(strcmp(opt[0], "mkcurve" ) == 0)
		return CREATE_CURVE;
	else if(strcmp(opt[0], "mkboard" ) == 0)
		return CREATE_CHARUCO;
	else if(strcmp(opt[0], "show" ) == 0)
		return SHOW_IMAGE;
	else if(strcmp(opt[0], "exit" ) == 0)
		return EXIT;
	return -1;
}

cv::Mat openImageImg(char* fileName)
{
	int fd = open(fileName, O_RDONLY);
	if(fd < 0)
	{
		Log(Log::WARN)<<"could not open "<<fileName<< "ignoreing";
		return cv::Mat();
	}
	
	size_t pos = 0;
	std::vector<char> buffer(IMREAD_SIZE);
	size_t count;
	while((count = read(fd, buffer.data()+pos, IMREAD_SIZE)) > 0)
	{
		pos+=count;
		buffer.resize(pos+IMREAD_SIZE);
		Log(Log::WARN)<<pos<<"  "<<IMREAD_SIZE;
	}
	close(fd);
	return cv::imdecode(buffer, cv::IMREAD_UNCHANGED);
}

cv::Mat openImageYml(char* fileName)
{
	cv::Mat image;
	cv::FileStorage matf(fileName, cv::FileStorage::READ);
	matf["image"]>>image;

	if(matf.isOpened() && !image.data)
	{
		Log(Log::WARN)<<fileName<<" dose not contain a valid image";
		matf.release();
		return cv::Mat();
	}
	else if(!image.data)
	{
		Log(Log::WARN)<<"could not open "<<fileName<< "ignoreing";
		matf.release();
		return cv::Mat();
	}
	return image;
}

std::vector<cv::Mat> loadImages(char** fileNames)
{ 
	std::vector<cv::Mat> images;
	for(size_t i = 0; fileNames[i]; ++i)
	{
		cv::Mat tmpImage;
		const std::string str(fileNames[i]);
		if(str.find(".mat") !=  std::string::npos)
		{
			Log(Log::DEBUG)<<__func__<<": "<<fileNames[i]<<" as YAML image";
			tmpImage = openImageYml(fileNames[i]);
		}
		else
		{
			Log(Log::DEBUG)<<__func__<<": "<<fileNames[i]<<" as png image";
			tmpImage = openImageImg(fileNames[i]);
		}
		if(tmpImage.data)
			images.push_back(tmpImage);
		else 
			Log(Log::WARN)<<"can not read image "<<i<<" from "<<fileNames[i];
	}
	return images;
}

int perfromOperation(int operation, char** fileNames, const Config& config)
{
	std::vector<cv::Mat> inImages;
	if(operation == CREATE_MAP || operation == APPLY_MAP || operation == APPLY_CURVE || operation == SHOW_IMAGE)
	{
		inImages = loadImages(fileNames);
		
		if(inImages.empty()) 
		{
			Log(Log::ERROR)<<"Input images must be provided";
			return -1;
		}
	}
	
	if(operation == CREATE_CHARUCO)
	{
		std::string fileName = config.output.empty() ? "out.png" : config.output;
		createCharucoBoard(config.size*X_BOARD_SIZE, fileName);
		Log(Log::INFO)<<"Exported charuco map of size "<<config.size*14<<" to "<<fileName;
		return 0;
	}
	else if(operation == CREATE_MAP)
	{
		cv::Mat mask;
		if(config.verbose)
		{
			cv::imshow( "Viewer", inImages[0] );
			cv::waitKey(0);
		}
		if(!config.bg.empty())
		{
			cv::Mat bg = cv::imread(config.bg);
		
			if(bg.data)
			{
				createMask(inImages[0], mask, bg);
				if(config.verbose)
				{
					cv::Mat masked;
					inImages[0].copyTo(masked, mask);
					cv::imshow( "Viewer", masked );
					cv::waitKey(0);
				}
			}
			else Log(Log::WARN)<<"can not read background image from "<<config.bg;
		}
		std::vector<DetectedPoint> points;
		
		if(config.harris)
			points = harrisDetectPoints(inImages[0], mask);
		else
			points = detectCharucoPoints(inImages[0], config.verbose);
		
		Log(Log::INFO)<<"Found "<<points.size()<<" points";
		if(points.size() < 8) 
		{
			Log(Log::ERROR)<<"Error creating map, insufficant points detected";
			return -1;
		}
		
		RemapMap map;
		map.outputCellSize = 50;
		createRemapMap(inImages[0], map, points, config.verbose);
		saveRemapMap(map, std::string(fileNames[0]) + ".mat");
	}
	else if(operation == APPLY_MAP)
	{
		std::vector<RemapedImage> remapedImages;
		
		for(size_t i = 0; i<inImages.size(); ++i)
		{
			Log(Log::INFO)<<"Processing image "<<i;
			RemapMap map = loadRemapMap(std::string(fileNames[i])+".mat");
			map.outputCellSize = config.size;
			
			if(!map.xMat.data)
			{
				Log(Log::ERROR)<<"could not load remap map from "<<std::string(fileNames[i])+".mat";
				return -1;
			}
			
			RemapedImage norm;
			if(!config.norm.empty())
			{
				cv::Mat tmp = cv::imread(config.norm);
				if(!tmp.data)
				{
					Log(Log::WARN)<<"could not open normalize file "<<config.norm;
				}
				norm = applyRemap(tmp, map);
				if(config.verbose)
				{
					cv::imshow("Viewer", norm.image);
					cv::waitKey(0);
				}
			}
			
			Log(Log::INFO)<<"Remaping image";
			RemapedImage remaped = applyRemap(inImages[i], map);
			
			Log(Log::INFO)<<"Normalizeing image";
			if(norm.image.data) normalize(remaped.image, norm.image);
			
			if(config.verbose)
			{
				cv::imshow( "Viewer", remaped.image );
				cv::waitKey(0);
				cv::imshow( "Viewer", remaped.angle);
				cv::waitKey(0);
			}
			
			applyKfactor(remaped.image, remaped.angle, config.kFactor);
			if(config.verbose)
			{
				cv::imshow("Viewer", remaped.image );
				cv::waitKey(0);
			}
			
			remapedImages.push_back(remaped);
		}
			
		cv::Mat out;
		
		Log(Log::INFO)<<"Stiching images";
		if(config.simpleStich)
			out = simpleStich(remapedImages);
		else 
			out = stich(remapedImages, true);
		
		if(config.verbose)
		{
			cv::imshow( "Viewer", out );
			cv::waitKey(0);
		}
		
		cv::imwrite(!config.output.empty() ? config.output : "out.png", out);
	}
	else if(operation == APPLY_CURVE)
	{
		if(config.curve.empty())
		{
			Log(Log::INFO)<<"a curve must be supplied";
			return -1;
		}
		
		cv::FileStorage fs(config.curve, cv::FileStorage::READ);
		cv::Mat curve;
		fs["cal"]>>curve;
		if(!curve.data || curve.type() != CV_32FC1 || curve.rows != 2 || curve.cols < 3)
		{
			Log(Log::INFO)<<"invalid curve";
			return -1;
		}
		
		if(inImages[0].channels() > 1)
			cvtColor(inImages[0], inImages[0], cv::COLOR_BGR2GRAY);
		if(inImages[0].type() != CV_32FC1)
		inImages[0].convertTo(inImages[0], CV_32F);
		
		Log(Log::DEBUG)<<"applyCurve";
		applyCurve(inImages[0], curve);
		cv::FileStorage fsO("out.mat", cv::FileStorage::WRITE);
		fsO<<"image"<<inImages[0];
		fsO.release();
	}
	else if(operation == CREATE_CURVE)
	{
		std::cout<<"how many coordinate pares are required?\n> ";
		int num = 0;
		std::cin>>num;
		if(std::cin.fail())
		{
			std::cin.clear();
			std::cout<<"invalid number";
			return -1;
		}
		cv::Mat curve = cv::Mat::zeros(2, num, CV_32FC1);
		float* keys   = curve.ptr<float>(0);
		float* values = curve.ptr<float>(1);
		std::cout<<"Please type "<<num<<" coordinate pairs\n";
		for(int i = 0; i < curve.cols; ++i)
		{
			std::cout<<i<<"> ";
			double key;
			double value;
			
			std::cin>>key;
			if(std::cin.fail())
			{
				std::cin.clear();
				--i;
				continue;
			}
			std::cin>>value;
			if(std::cin.fail())
			{
				std::cin.clear();
				--i;
				continue;
			}
			
			
			keys[i] = key;
			values[i] = value;
		}
		
		for(int i = 0; i < curve.cols; ++i)
			std::cout<<keys[i]<<' '<<values[i]<<'\n';
		
		cv::FileStorage fs(!config.output.empty() ? config.output : "curve.mat", cv::FileStorage::WRITE);
		fs<<"cal"<<curve;
		fs.release();
		
		std::cout<<"Curve saved to "<<(!config.output.empty() ? config.output : "curve.mat")<<'\n';
	}
	else if(operation == SHOW_IMAGE)
	{
		cv::namedWindow("Show Image", cv::WINDOW_NORMAL );
		for(size_t i = 0; i < inImages.size(); ++i)
		{
			cv::imshow("Show Image", inImages[i]);
			cv::waitKey(0);
		}
		cv::destroyWindow("Show Image"); 
	}
	else if(operation == EXIT)
		return -1;
	return 0;
}

int main(int argc, char* argv[])
{
	cv::ocl::setUseOpenCL(false);
	
	Config config;
    argp_parse(&argp, argc, argv, 0, 0, &config);
	
	Log::level = config.quiet ? Log::WARN : config.verbose ? Log::DEBUG : Log::INFO;
	
	Log(Log::INFO)<<"UVOS Optical lubricant thikness mapper "<<argp_program_version;

	if(config.verbose) 
	{
		cv::namedWindow( "Viewer", cv::WINDOW_NORMAL );
		cv::resizeWindow("Viewer", 960, 500);
	}
	
	int operation = selectOperation(config.commandsFiles);
	
	int ret = 0;
	
	if(!config.interactive)
	{
		if(operation < 0)
		{
			Log(Log::ERROR)<<"An operation must be selected";
			Log(Log::INFO)<<"Possible operations: apply create curve mkcurve mkboard";
			return -1;
		}
		ret = perfromOperation(operation, config.commandsFiles+1, config);
	}
	else
	{
		while(true)
		{
			std::cout<<"> ";
			char cmdline[1024];
			char** tokens = new char*[256];
			char *token;
			char *r = cmdline;
			
			for(size_t i = 0; i < 256/sizeof(char*); ++i) 
				tokens[i] = nullptr;
			
			std::cin.getline(cmdline, sizeof(cmdline), '\n');
			
			for(size_t i = 0; (token = strtok_r(r, " ", &r)) && i < 256; ++i)
			{
				tokens[i] = new char[strlen(token)+1];
				strcpy(tokens[i], token);
			}
			
			int operation = selectOperation(tokens);
			if(operation < 0)
			{
				Log(Log::ERROR)<<"A operation must be selected";
				continue;
			}
			
			if(perfromOperation(operation, tokens+1, config) < 0)
				break;
			
			for(size_t i = 0; i < 256/sizeof(char*) && tokens[i]; ++i) 
				delete[] tokens[i];
			delete[] tokens;
		}
	}
	
	if(config.verbose)
	{
		cv::destroyWindow("Viewer");  
		cv::waitKey(0);
	}
	
	return ret;
}