//neural network/ prediction
#include <iostream>
#include "opencv2/opencv.hpp"
#include <stdio.h>
#include "opencv2/imgproc/imgproc.hpp"
#include "opencv2/highgui/highgui.hpp"
#include <stdlib.h>
#include <opencv2/core/core.hpp>
#include <opencv2/ml/ml.hpp>
#include <limits>

using namespace cv;
using namespace std;
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
Mat loadfile(string& ymlname) //load a file from .yml archive
{
	Mat featforTraining;
	cv::FileStorage forTraining(ymlname, cv::FileStorage::READ);
	forTraining["features"] >> featforTraining;
	forTraining.release();
	return featforTraining;
}
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
Mat scriviperriga(cv::Mat& ref, int& w) //write a matrix in a new matrix of 1 row
{
	int n_samples = (ref.rows-w+1)*(ref.cols-w+1);

	Mat effettiva(ref.rows-w+1,ref.cols-w+1,CV_32FC1);
	Mat refperriga(1,n_samples,CV_32FC1);

	int c=0;
	for(int i=0; i<ref.cols-w+1;i++)
	{
		for(int j=0; j<ref.rows-w+1;j++)
		{
			refperriga.at<float>(0,c) = ref.at<float>(i+floor(w/2),j+floor(w/2));
			c++;
		}
	}
	Mat rigatrasp;
	transpose(refperriga,rigatrasp);
	return rigatrasp;
}
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
Mat setlabels(cv::Mat& riga, int& n_classes, int& nero, int& bianco) //set labels for learning
{
	Mat labelsforTraining(riga.rows,n_classes,CV_32FC1);
	for(int j=0; j<riga.rows; j++)
	{
		if (riga.at<float>(j,0)==0)
		{
			labelsforTraining.at<float>(j,0)=nero;
		}
		else
		{
			labelsforTraining.at<float>(j,0)=bianco;
		}
	}
	return labelsforTraining;
}
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//method that does prediction and produces final image
//crossvalidation of data in learning is not implemented

Mat predictproduce(CvANN_MLP& network, cv::Mat& feattoPredict, int& n_features, int& w, int& n_classes, float& thresh)
{
	Mat test_sample(1,n_features,CV_32FC1);
	Mat classificationResult(1,n_classes,CV_32FC1);

	int outputdimension = sqrt(feattoPredict.rows);
	Mat image_result_riga(1,feattoPredict.rows,CV_32FC1);
	Mat image_result(outputdimension,outputdimension,CV_32FC1);

	int count=1;
	for(int i=0; i<feattoPredict.rows; i++)
	{
		for(int j=0; j<feattoPredict.cols; j++) //get a row
		{
			test_sample.at<float>(0,j)=feattoPredict.at<float>(i,j);
		}
		network.predict(test_sample,classificationResult);
		int pixel;
		if (1-abs(classificationResult.at<float>(0,0))>=thresh) pixel=255; else pixel=0;
		image_result_riga.at<float>(0,count-1)=pixel; //salvo l'immagine per riga
		count++;
	}

	//write final I_CO image
	int count2 = 0;
	for(int i=0; i<image_result.cols; i++)
	{
		for(int j=0; j<image_result.rows; j++)
		{
			image_result.at<float>(i,j)=image_result_riga.at<float>(0,count2);
			count2++;
		}
	}
	return image_result;
}
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//method for training
void training(CvANN_MLP& network, CvANN_MLP_TrainParams& params, int& nero, int& bianco, int& w, int& n_classes, cv::Mat& featforTraining, cv::Mat& ref)
{
	Mat riga = scriviperriga(ref,w);
	Mat labelsforTraining = setlabels(riga,n_classes,nero,bianco);
	network.train(featforTraining,labelsforTraining,Mat(),Mat(),params);
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

int main( int argc, char** argv )
{
	//parameters
	int n_features = 7, w = 17, n_classes = 1, nero = 1, bianco = 0;
	float thresh = 0.45;
	string yml,tif;

	//list of strings for recursive learning/predicting over DRIVE database

	string listtraintotal[] = {"01.yml","02.yml","03.yml","04.yml","05.yml","06.yml","07.yml","08.yml","09.yml","10.yml",
							   "11.yml","12.yml","13.yml","14.yml","15.yml","16.yml","17.yml","18.yml","19.yml","20.yml",
							   "21.yml","22.yml","23.yml","24.yml","25.yml","26.yml","27.yml","28.yml","29.yml","30.yml",
							   "31.yml","32.yml","33.yml","34.yml","35.yml","36.yml","37.yml","38.yml","39.yml","40.yml"};

	string listreftotal[] = {"01_manual.tif","02_manual.tif","03_manual.tif","04_manual.tif","05_manual.tif","06_manual.tif","07_manual.tif","08_manual.tif","09_manual.tif","10_manual.tif",
							 "11_manual.tif","12_manual.tif","13_manual.tif","14_manual.tif","15_manual.tif","16_manual.tif","17_manual.tif","18_manual.tif","19_manual.tif","20_manual.tif",
							 "21_manual.tif","22_manual.tif","23_manual.tif","24_manual.tif","25_manual.tif","26_manual.tif","27_manual.tif","28_manual.tif","29_manual.tif","30_manual.tif",
							 "31_manual.tif","32_manual.tif","33_manual.tif","34_manual.tif","35_manual.tif","36_manual.tif","37_manual.tif","38_manual.tif","39_manual.tif","40_manual.tif"};

	string listtrain[] = {"01.yml","02.yml","03.yml","04.yml","05.yml","06.yml","07.yml","08.yml","09.yml","10.yml",
						"11.yml","12.yml","13.yml","14.yml","15.yml","16.yml","17.yml","18.yml","19.yml","20.yml"};

	string listtrain2[] = {"21.yml","22.yml","23.yml","24.yml","25.yml","26.yml","27.yml","28.yml","29.yml","30.yml",
						 "31.yml","32.yml","33.yml","34.yml","35.yml","36.yml","37.yml","38.yml","39.yml","40.yml",};

	string listref[] = {"01_manual.tif","02_manual.tif","03_manual.tif","04_manual.tif","05_manual.tif","06_manual.tif","07_manual.tif","08_manual.tif","09_manual.tif","10_manual.tif",
						"11_manual.tif","12_manual.tif","13_manual.tif","14_manual.tif","15_manual.tif","16_manual.tif","17_manual.tif","18_manual.tif","19_manual.tif","20_manual.tif"};

	string listref2[] = {"21_manual.tif","22_manual.tif","23_manual.tif","24_manual.tif","25_manual.tif","26_manual.tif","27_manual.tif","28_manual.tif","29_manual.tif","30_manual.tif",
						 "31_manual.tif","32_manual.tif","33_manual.tif","34_manual.tif","35_manual.tif","36_manual.tif","37_manual.tif","38_manual.tif","39_manual.tif","40_manual.tif"};

	string listpred[] = {"01.yml","02.yml","03.yml","04.yml","05.yml","06.yml","07.yml","08.yml","09.yml","10.yml",
						 "11.yml","12.yml","13.yml","14.yml","15.yml","16.yml","17.yml","18.yml","19.yml","20.yml",
						 "21.yml","22.yml","23.yml","24.yml","25.yml","26.yml","27.yml","28.yml","29.yml","30.yml",
						 "31.yml","32.yml","33.yml","34.yml","35.yml","36.yml","37.yml","38.yml","39.yml","40.yml",};

	string listout[] = {"01p.tif","02p.tif","03p.tif","04p.tif","05p.tif","06p.tif","07p.tif","08p.tif","09p.tif","10p.tif",
						"11p.tif","12p.tif","13p.tif","14p.tif","15p.tif","16p.tif","17p.tif","18p.tif","19p.tif","20p.tif",
						"21p.tif","22p.tif","23p.tif","24p.tif","25p.tif","26p.tif","27p.tif","28p.tif","29p.tif","30p.tif",
						"31p.tif","32p.tif","33p.tif","34p.tif","35p.tif","36p.tif","37p.tif","38p.tif","39p.tif","40p.tif",};

	//-----------------------------------------------------NEURAL NETWORK
	int layers_d[] = {n_features,15,15,15,n_classes};
	Mat layers(1,sizeof(layers_d)/4,CV_32SC1,layers_d);

	CvANN_MLP network;
	network.create(layers,CvANN_MLP::SIGMOID_SYM,1,1);
	CvANN_MLP_TrainParams params = CvANN_MLP_TrainParams(
								cvTermCriteria(CV_TERMCRIT_ITER + CV_TERMCRIT_EPS, 1e9, 1e-6),
								CvANN_MLP_TrainParams::BACKPROP, 0.0000001, 0.0000001);

	//-------------------------------------------------------------TRAINING
	int count=1;

	Mat featforTraining, ref;

	for(int i=0; i<sizeof(listtrain)/4;i++)
	{
		string inputyml = listtrain[i];
		string reftif = listref[i];

		featforTraining = loadfile(inputyml);
		ref = imread(reftif,0); ref.convertTo(ref,CV_32F);
		training(network,params,nero,bianco,w,n_classes,featforTraining,ref);

		cout<<count<<endl;
		count++;
	}

	//---------------------------------------------------------------------TESTING DATA

	for(int i=0; i<sizeof(listpred)/4;i++)
	{
		string archive = listpred[i];
		string outputname = listout[i];

		Mat feattoPredict = loadfile(archive);
		Mat image_result = predictproduce(network,feattoPredict,n_features,w,n_classes,thresh);
		imwrite(outputname,image_result);
	}

	waitKey(0);
	return 0;
}




