#include "hCRF.h"
#include <iostream>
#include <string>
#ifdef WIN32
#include <conio.h>
#endif

#include <time.h> 

#ifdef _OPENMP
#include <omp.h>
#endif

using namespace std;

#define MODE_TRAIN 1
#define MODE_TEST  2
#define MODE_DEBUG 4
 
#define TOOLBOX_HCRF			1
#define TOOLBOX_LDCRF			2 
#define TOOLBOX_MVHCRF			4
#define TOOLBOX_MVLDCRF			8 
 

void usage (char **argv)
{  
  cerr << "HCRF library"<< endl << endl;
  cerr << "usage> " << argv[0] << " [-d filename] [-l filename] [-sl filename] [-D filename] [-L filename] [-SL filename] [-m filename] [-f filename] [-r filename] [-o cg|bfgs]" << endl << endl;
  cerr << "options:" << endl;
  cerr << " -h\tNumber of hidden states (def.=3)" << endl;  
  cerr << " -a\tSelect the model type (hcrf, ldcrf, mvhcrf, mvldcrf) (def. = mvhcrf)" << endl;
  cerr << " -e\tSelect the graph type for multiview models (linked, coupled, linkedcoupled) (def. = linked)" << endl;
  cerr << " -b\tSelect the inference engine (fb, jt, lbp) (def. = jt)" << endl;
  cerr << " -o\tSelect optimizer: lbfgs or nrbm (def.= lbfgs)" << endl;
  cerr << " -d\tName of file containing the training data (def.= dataTrain.csv)" << endl;
  cerr << " -l\tName of file containing the training labels (def.= labelsTrain.csv)" << endl;
  cerr << " -q\tName of file containing the training sequence labels (def.= seqLabelsTrain.csv)" << endl;
  cerr << " -D\tName of file containing the testing data (def.= dataTest.csv)" << endl;
  cerr << " -L\tName of file containing the testing labels (def.= labelsTest.csv)" << endl;
  cerr << " -Q\tName of file containing the testing sequence labels (def.= seqLabelsTest.csv)" << endl;
  cerr << " -m\tName of file where model is written (def.= model.txt)" << endl;
  cerr << " -f\tName of file where the features are written (def.= features.txt)" << endl;
  cerr << " -F\tName of file where the features types are written (def.= none, optional)" << endl;
  cerr << " -r\tName of file where the computed labels are written (def.= results.txt)" << endl;
  cerr << " -c\tName of file where the statistics are written (def.= stats.txt)" << endl;
  cerr << " -i\tMaximum number of iterations (def.= 300)" << endl;
  cerr << " -s2\tSigma2: L2-norm regularization factor(def.= 0.0)" << endl;
  cerr << " -p\tDebug print level (def.= 1)" << endl;
  cerr << endl;
  exit(1);
}

int main(int argc, char **argv)
{   
	int toolboxType, opt, maxIt, nbHiddenStates, debugLevel;
	double regFactorL2;

	Toolbox* toolbox = NULL;
	toolboxType = TOOLBOX_MVHCRF;
	InferenceEngine* engine = NULL;
	eGraphTypes graphType = CHAIN;
	opt = OPTIMIZER_LBFGS;
	maxIt = 300;  
	nbHiddenStates = 3;
	debugLevel = 1;
	regFactorL2 = 0.0;
	
	long seed = 0;

	// For the multiview models
	int nbViews = 2;
	std::vector<int> nbHiddenStatesMV;
	nbHiddenStatesMV.push_back(3);
	nbHiddenStatesMV.push_back(3);
	std::vector<int>  rfi_v1, rfi_v2;
	for(int i=0; i<12; i++)  rfi_v1.push_back(i);
	for(int i=12; i<20; i++)  rfi_v2.push_back(i);
	std::vector<std::vector<int>> rawFeatureIndex;
	rawFeatureIndex.push_back(rfi_v1);
	rawFeatureIndex.push_back(rfi_v2); 

#ifdef UNIX
	string home = "./";
#else
	string home = ".\\";
#endif
	string filenameDataTrain = home + "dataTrain.csv";
	string filenameLabelsTrain = home + "labelsTrain.csv" ;
	string filenameSeqLabelsTrain = home + "seqLabelsTrain.csv";
	string filenameDataTest = home + "dataTest.csv";
	string filenameLabelsTest = home + "labelsTest.csv";
	string filenameSeqLabelsTest = home + "seqLabelsTest.csv";
	string filenameModel = home + "model.txt";
	string filenameFeatures = home + "features.txt";
	string filenameOutput = home + "results.txt";
	string filenameStats = home + "stats.txt";

	/* Read command-line arguments */
	for (int k=1; k<argc; k++)
	{
		if(argv[k][0] != '-') 
			break;
		else if(argv[k][1] == 'd')  
			filenameDataTrain = argv[++k];
		else if(argv[k][1] == 'l') 
			filenameLabelsTrain = argv[++k];
		else if(argv[k][1] == 'q') 
			filenameSeqLabelsTrain = argv[++k];
		else if(argv[k][1] == 'D') 
			filenameDataTest = argv[++k];
		else if(argv[k][1] == 'L') 
			filenameLabelsTest = argv[++k];
		else if(argv[k][1] == 'Q') 
			filenameSeqLabelsTest = argv[++k];
		else if(argv[k][1] == 'm') 
			filenameModel = argv[++k];
		else if(argv[k][1] == 'f') 
			filenameFeatures = argv[++k];
		else if(argv[k][1] == 'r') 
			filenameOutput = argv[++k];
		else if(argv[k][1] == 'c') 
			filenameStats = argv[++k];
		else if(argv[k][1] == 'p')
			debugLevel = atoi(argv[++k]);
		else if(argv[k][1] == 'i')
			maxIt = atoi(argv[++k]);
		else if(argv[k][1] == 'h')
			nbHiddenStates = atoi(argv[++k]);
		else if(argv[k][1] == 'Z') 
			seed = atol(argv[++k]);
		else if(argv[k][1] == 's')
			regFactorL2 = atof(argv[++k]);
		else if(argv[k][1] == 'o') {
			if(!strcmp(argv[k+1],"lbfgs"))
				opt = OPTIMIZER_LBFGS;
			else if(!strcmp(argv[k+1],"nrbm")) {
				opt = OPTIMIZER_NRBM; 
			}
			k++;
		} 
		else if(argv[k][1] == 'a') {
			if(!strcmp(argv[k+1],"hcrf"))
				toolboxType = TOOLBOX_HCRF;
			else if(!strcmp(argv[k+1],"ldcrf"))
				toolboxType = TOOLBOX_LDCRF;
			else if(!strcmp(argv[k+1],"mvhcrf"))
				toolboxType = TOOLBOX_MVHCRF;
			else if(!strcmp(argv[k+1],"mvldcrf"))
				toolboxType = TOOLBOX_MVLDCRF;
			k++;
		}
		else if(argv[k][1] == 'b') { 
			if(!strcmp(argv[k+1],"fb"))
				engine = new InferenceEngineFB();
			else if(!strcmp(argv[k+1],"jt"))
				engine = new InferenceEngineJT();
			else if(!strcmp(argv[k+1],"lbp"))
				engine = new InferenceEngineLoopyBP();
			k++;
		}
		else if(argv[k][1] == 'e') {
			if(!strcmp(argv[k+1],"linked"))
				graphType = MV_GRAPH_LINKED;
			if(!strcmp(argv[k+1],"coupled"))
				graphType = MV_GRAPH_COUPLED;
			if(!strcmp(argv[k+1],"linkedcoupled"))
				graphType = MV_GRAPH_LINKED_COUPLED;
			k++;
		}
		else usage(argv);
    }

	if(toolboxType == TOOLBOX_HCRF)
		toolbox = new ToolboxHCRF(nbHiddenStates);
	else if(toolboxType == TOOLBOX_LDCRF)
		toolbox = new ToolboxLDCRF(nbHiddenStates); 
	else if(toolboxType == TOOLBOX_MVHCRF) 
		toolbox = new ToolboxMVHCRF(graphType, nbViews, nbHiddenStatesMV, rawFeatureIndex);
	else if(toolboxType == TOOLBOX_MVLDCRF)
		toolbox = new ToolboxMVLDCRF(graphType, nbViews, nbHiddenStatesMV, rawFeatureIndex); 

	toolbox->setOptimizer(opt); 
	toolbox->setInferenceEngine(engine);
	toolbox->setDebugLevel(debugLevel);
	toolbox->setRandomSeed(seed);  

	// Train
	{
		cout << "Reading training set..." << endl;
		DataSet data;
		const char* fileData = filenameDataTrain.empty()? 0 : filenameDataTrain.c_str();
		if(toolboxType == TOOLBOX_HCRF || toolboxType == TOOLBOX_MVHCRF )
			data.load(fileData,NULL,(char*)filenameSeqLabelsTrain.c_str(),NULL,NULL);
		else
			data.load(fileData,(char*)filenameLabelsTrain.c_str(),NULL,NULL,NULL);

		toolbox->setMaxNbIteration(maxIt);
		toolbox->setRegularizationL2(regFactorL2);
		toolbox->setRangeWeights(-1,1);
		toolbox->setWeightInitType(INIT_RANDOM);

		cout << "Starting training ..." << endl;
		toolbox->train(data,true);
		toolbox->save((char*)filenameModel.c_str(),(char*)filenameFeatures.c_str());
	}

	// Test
	{
		cout << "Reading testing set..." << endl;
		DataSet data;
		if(toolboxType == TOOLBOX_HCRF || toolboxType == TOOLBOX_MVHCRF )
			data.load((char*)filenameDataTest.c_str(),NULL,(char*)filenameSeqLabelsTest.c_str());
		else
			data.load((char*)filenameDataTest.c_str(),(char*)filenameLabelsTest.c_str());

		ofstream fileStats1 ((char*)filenameStats.c_str()); 
		if (fileStats1.is_open())
		{
			fileStats1 << endl << endl << "TESTING DATA SET" << endl << endl;
			fileStats1.close();
		}

		cout << "Starting testing ..." << endl;
		toolbox->load((char*)filenameModel.c_str(),(char*)filenameFeatures.c_str());
		toolbox->initToolbox();
		toolbox->test(data,(char*)filenameOutput.c_str(),(char*)filenameStats.c_str()); 
	}

	if(toolbox)
		delete toolbox;
	
	cout << "Press a key to continue..." << endl;
#ifdef WIN32
	_getch();
#endif

	return 0;
} 

