// -*- Mode: C++; tab-width: 2; -*-
// vi: set ts=2:
//

#include <BALL/FORMAT/commandlineParser.h>
#include <BALL/QSAR/registry.h>
#include <BALL/QSAR/featureSelection.h>
#include <BALL/QSAR/configIO.h>
#include <BALL/QSAR/registry.h>
#include <fstream>

using namespace BALL::QSAR;
using namespace BALL;
using namespace std;

void startModelCreation(ModelConfiguration& conf, QSARData* q, String* data_filename);


void startModelCreation(ifstream& in, QSARData* q, String* data_filename)
{
	ModelConfiguration conf = ConfigIO::readModelConfiguration(&in);
	if(conf.done) return; // stop processing this section

	startModelCreation(conf, q, data_filename);
}


void startModelCreation(ModelConfiguration& conf, QSARData* q, String* data_filename)
{
	bool created_data_object=0;
	if(q==NULL || data_filename==NULL || conf.data_file!=*data_filename)
	{
		if(q==NULL)
		{
			q = new QSARData;
			created_data_object=1;
		}
		q->readFromFile(conf.data_file);
		if(data_filename) *data_filename = conf.data_file;
	}
	else
	{
		Log.level(2)<<"[ModelCreator debug-info:] QSARData object for file "<<conf.data_file<<" already in memory; not reading it again."<<endl;
	}
	Registry reg;
	Model* model;
	bool kernel=0;
	RegistryEntry* entry = reg.getEntry(conf.model_no);

	if(entry->create!=NULL)
	{
		model = (*entry->create)(*q);
	}
	else
	{
		if(conf.kernel_type==0 || conf.kernel_par1==0)
		{
			Log.error()<<"For kernel based model, kernel-type and kernel-parameter(s) must be specified!"<<endl;
			return;
		}

		model = (*entry->createKernel1)(*q,conf.kernel_type,conf.kernel_par1, conf.kernel_par2);
		kernel=1;
	}

	if(conf.model_parameters.size()>0)
	{
		model->setParameters(conf.model_parameters);
	}
	if(!conf.no_training && conf.optimize_model_parameters)
	{
		if(conf.k_fold==0)
		{
			Log.error()<<"'k_fold' must be set if model parameters are to be optimized!"<<endl;
			return;
		}
		model->optimizeParameters(conf.k_fold);
	}
	if(!conf.no_training && kernel && conf.grid_search_steps>0)
	{
		if(conf.k_fold==0)
		{
			Log.error()<<"'k_fold' must be set if grid search is to be done!"<<endl;
			return;
		}
		if(conf.grid_search_stepwidth==0 && conf.kernel_type!=2)
		{
			Log.error()<<"'grid_search_stepwidth' must be set if grid search is to be done!"<<endl;
			return;
		}
		((KernelModel*)model)->kernel->gridSearch(conf.grid_search_stepwidth, conf.grid_search_steps, conf.grid_search_recursions, conf.k_fold);
	}

	model->readTrainingData();
	if(!conf.no_training)
	{
		try
		{
			model->train();
		}
		catch(BALL::Exception::GeneralException e)
		{
			Log.error()<<e.getMessage();
		}
	}

	model->saveToFile(conf.output);

	if(created_data_object) delete q;
	delete model;
}



#ifndef EXT_MAIN
int main(int argc, char* argv[])
{
	CommandlineParser par("ModelCreator","create a QSAR model       ","1.1",String(__DATE__), "QuEasy (QSAR)");
	par.registerMandatoryInputFile("i", "input dat-file");
	par.registerMandatoryOutputFile("o", "output model file");
	par.registerMandatoryStringParameter("type", "model type");
	par.registerOptionalStringParameter("kernel", "kernel type (in case of kernel-model)");
	Registry reg;
	list<String> restr;
	for(RegistryEntryIterator it=reg.beginEntry(); it!=reg.endEntry(); it++)
	{
		restr.push_back(it->second.name_abreviation);
	}
	par.setParameterRestrictions("type", restr);
	restr.clear();
	restr.push_back("none");
	restr.push_back("polynomial");
	restr.push_back("rbf");
	restr.push_back("sigmoidal");
	par.setParameterRestrictions("kernel", restr);
	String man = "ModelCreator creates a QSAR model using an input data set as generated by InputReader.\n\nThe type of QSAR model to be used can be specified by '-type', the type of kernel-function (if any) can be chosen by '-kernel'. Optimization of model- and kernel-parmeters will be done automatically using cross-validation.\n\nOutput of this tool is a model-file that can be used by other QuEasy tools (e.g. FeatureSelector).";
	par.setToolManual(man);
	par.setSupportedFormats("i","dat");
	par.setSupportedFormats("o","mod");
	par.parse(argc,argv);

	ModelConfiguration conf;
	conf.data_file = par.get("i");
	conf.output = par.get("o");
	conf.optimize_model_parameters = true;
	conf.kernel_par1 = reg.default_kernel_par1;
	conf.kernel_par1 = reg.default_kernel_par2;
	conf.k_fold = reg.default_k;
	conf.grid_search_recursions = reg.default_gridsearch_recursion;
	conf.grid_search_stepwidth = reg.default_gridsearch_stepwidth;
	conf.grid_search_steps = reg.default_gridsearch_steps;
	try
	{
		conf.model_no = reg.getModelNo(par.get("type"));
	}
	catch(BALL::Exception::GeneralException e)
	{
		cerr << "A model-type '"<<par.get("type")<<"' does not exist; possible choices are:"<<endl;
		for(RegistryEntryIterator it=reg.beginEntry(); it!=reg.endEntry(); it++)
		{
			cerr<<"   "<<it->second.name_abreviation<<" :  "<<it->second.name<<" ";
			if(it->second.regression) cerr<<"(regression)"<<endl;
			else cerr<<"(classification)"<<endl;
		}
		cerr<<endl;
		exit(1);
	}

	String kernel = par.get("kernel");
	if(kernel!=CommandlineParser::NOT_FOUND && kernel != "none")
	{
		if(!reg.getEntry(par.get("type"))->kernel)
		{
			cerr << "[Error:] The chosen model-type has no kernel but you specified a kernel-type!"<<endl;
			exit(1);
		}
		if(kernel=="polynomial")
		{
			conf.kernel_type = 1;
		}
		else if(kernel=="rbf")
		{
			conf.kernel_type = 2;
		}
		else if(kernel=="sigmoidal")
		{
			conf.kernel_type = 3;
		}
		else
		{
			cerr << "Specified kernel-type '"<<kernel<<"' unknown; possible choices are: polynomial, rbf, sigmoidal"<<endl;
			exit(1);
		}
	}

	startModelCreation(conf,0,0);

}
#endif
