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

using namespace BALL;
using namespace BALL::Docking;

void processMultiMoleculeFile(ScoringFunction* scoring_function, StructurePreparer* sp, String par_file, Rescoring* rescoring, bool train, double min_dock_score, String dock_score_label, GenericMolFile* input, GenericMolFile* output, double output_score_threshold, bool ignore_top)
{
	list<pair<double, bool> > rescore_list;
	double min_rescore = 1e12;
	int i = 1;

	for (Molecule* mol = input->read(); mol; mol = input->read(), i++)
	{
		String name = mol->getName();
		Log<<"====== Ligand "<<i;
		if (name != "") Log<<", "<<name;
		Log<<" ============"<<endl;
		Log.flush();

		if (mol->hasProperty("score_ligcheck"))
		{
			double score_ligcheck = ((String)mol->getProperty("score_ligcheck").toString()).toDouble();
			if (score_ligcheck < 0.95) // 0 = error, 1 = check passed
			{
				cout<<"Skipping compound because it has been marked as containing errors by LigCheck."<<endl;
				delete mol;
				continue;
			}
		}

		double score = 0;
		try
		{
			sp->prepare(mol, par_file);
			if (!rescoring || train)
			{
				scoring_function->setLigand(mol);
				scoring_function->update();
				score = scoring_function->updateScore();
				scoring_function->printResult();
			}
			if (rescoring)
			{
				if (train)
				{
					if (score > 1000)
					{
						Log.level(10)<<"score>1000, thus current compound is not added to training data set."<<endl;
						delete mol;
						continue;
					}
					rescoring->addScoreContributions(mol);
				}
				else score = rescoring->rescore(mol);
			}

			if (score < output_score_threshold)
			{
				if (ignore_top)
				{
					bool keep_dock_score = 0;
					if (mol->hasProperty("score"))
					{
						double docking_score = mol->getProperty(dock_score_label).toString().toDouble();
						if (docking_score < min_dock_score)
						{
							rescore_list.push_back(make_pair(docking_score, false));
							keep_dock_score = 1;
						}
					}
					if (!keep_dock_score)
					{
						rescore_list.push_back(make_pair(score, true));
						if (score < min_rescore)
						{
							min_rescore = score;
						}
					}
				}
				else
				{
					mol->setProperty("re-score", score);
				}
			}
		}
		catch (BALL::Exception::GeneralException e)
		{
			cout<<e.getMessage()<<endl;
			cout<<"Error! Skipping this molecule!"<<endl;
		}

		if (score < output_score_threshold && (!rescoring || !train) && !ignore_top)
		{
			*output << *mol;
		}

		delete mol;
	}

	/** If ignoring (i.e. not rescoring) the top fraction of docking-results, we need to make sure
	 that all rescored compounds get a rescore-value larger than that of the former. */
	if (ignore_top && !train)
	{
		input->reopen();
		int i = 1;
		double offset = min_dock_score-min_rescore+0.25;

		list<pair<double, bool> >::iterator rescore_it = rescore_list.begin();
		for (Molecule* mol = input->read(); mol; mol = input->read(), i++)
		{
			String name = mol->getName();
			if (mol->hasProperty("score_ligcheck"))
			{
				double score_ligcheck = ((String)mol->getProperty("score_ligcheck").toString()).toDouble();
				if (score_ligcheck < 0.95) // 0 = error, 1 = check passed
				{
					delete mol;
					continue;
				}
			}
			double rescore = rescore_it->first;
			if (rescore_it->second) // if compound was rescored
			{
				rescore += offset;
			}
			if (rescore < output_score_threshold)
			{
				mol->setProperty("re-score", rescore);
				*output << *mol;
			}
			delete mol;
			rescore_it++;
		}
	}

	delete input;
	if (output) delete output;
}


int runRescoring(CommandlineParser& par, bool simple_rescoring, bool train)
{
	// just to make sure ...
	if (simple_rescoring) train = false;


	/** If desired, write ini-file with default parameters and abort */

	String default_inifile = par.get("write_ini");
	if (default_inifile != CommandlineParser::NOT_FOUND)
	{
		// if ini-file already exists, read its entries first
		Options default_options;
		ScoringFunction::getDefaultOptions(default_options);
		list<Constraint*> clist;
		if (ifstream(default_inifile.c_str()))
		{
			DockingAlgorithm::readOptionFile(default_inifile, default_options, clist);
		}

		Options* scoring_options = default_options.getSubcategory(ScoringFunction::SUBCATEGORY_NAME);
		scoring_options->setDefault("scoring_type", "MM");
		scoring_options->setDefault("nonbonded_cutoff_precalculation", scoring_options->get("nonbonded_cutoff"));
		scoring_options->set("nonbonded_cutoff", 3.0);
		DockingAlgorithm::writeOptionFile(par.get("write_ini"), default_options, clist);
		Log << "Ini-file w/ default values has been written to file '"<<default_inifile<<"'. Goodbye!"<<endl;
		return 0;
	}


	/** Fetch information about specification of desired rescoring approach */

	String method = "";
	String free_energy_label = "";
	String model_file = "";
	String scoring_type = "MM";
	String grid_file = "";
	Options option;
	par.copyAdvancedParametersToOptions(option);
	Options* option_category = option.getSubcategory(ScoringFunction::SUBCATEGORY_NAME);
	if (!option_category) option_category = &option;

	list<Constraint*> constraints;
	if (par.get(DockingAlgorithm::OPTION_FILE_PARAMETER_NAME) != CommandlineParser::NOT_FOUND)
	{
		DockingAlgorithm::readOptionFile(par.get(DockingAlgorithm::OPTION_FILE_PARAMETER_NAME), option, constraints);
		scoring_type = option_category->setDefault("scoring_function", "MM");
		grid_file = option_category->setDefault("grid_file", "");
	}
	// Overload ini-file settings with values taken from command-line (if any)
	if (par.get("function") != CommandlineParser::NOT_FOUND)
	{
		scoring_type = par.get("function");
	}

	method = par.get("method");
	free_energy_label = par.get("exp");
	if (par.get("mod") != CommandlineParser::NOT_FOUND)
	{
		model_file = par.get("mod");
	}

	if (!simple_rescoring && method != "Rescoring3D" && method != "Rescoring4D" && method != "Rescoring1D")
	{
		cerr << "[Error:] Type of desired rescoring method unknown: available are 'Rescoring3D', 'Rescoring4D' and 'Rescoring1D'." << endl;
		exit(1);
	}


	/** Setup StructurePreparer, Rescoring and ScoringFunction  */

	//Log.setMinLevel(cout, 39);

	GenericMolFile* ref_ligand_file = MolFileFactory::open(par.get("rl"));
	Molecule* ref_ligand = ref_ligand_file->read();
	ref_ligand_file->close();
	delete ref_ligand_file;

	StructurePreparer* sp;
	if (scoring_type.hasSubstring("PLP"))
	{
		sp = new StructurePreparer("PLP");
	}
	else
	{
		sp = new StructurePreparer;
	}

	String par_file = option_category->get("filename");
	if (par_file == "") par_file="Amber/amber96-docking.ini";
	System receptor;
	GenericMolFile* receptor_file = MolFileFactory::open(par.get("rec"));
	if (!receptor_file)
	{
		cerr<<"Format of receptor-file not supported!"<<endl;
		return 1;
	}

	*receptor_file >> receptor;
	delete receptor_file;
	sp->prepare(&receptor, par_file);
	sp->prepare(ref_ligand, par_file);

	ScoringFunction* scoring_function;
	if (scoring_type == "MM" || scoring_type == "PB")
	{
		// Support for using one and the same config-file for docking and rescoring
		String precalc_nonbonded_cuttoff = option_category->get("nonbonded_cutoff_precalculation");
		if (precalc_nonbonded_cuttoff != "")
		{
			option_category->set("nonbonded_cutoff", precalc_nonbonded_cuttoff);
		}
	}
	if (scoring_type == "MM")
	{
		scoring_function = new MMScoring(receptor, *ref_ligand, option);
	}
	else if (scoring_type == "GridedMM")
	{
		scoring_function = new GridedMM(receptor, *ref_ligand, option);
	}
	else if (scoring_type == "GridedPLP")
	{
		scoring_function = new GridedPLP(receptor, *ref_ligand, option);
	}
	else if (scoring_type == "PLP")
	{
		scoring_function = new PLPScoring(receptor, *ref_ligand, option);
	}
	else if (scoring_type == "PB")
	{
		scoring_function = new PBScoring(receptor, *ref_ligand, option);
	}
	else
	{
		String mess="ScoringFunction type \'"+scoring_type+"\' unknown/unsupported!";
		cerr<<"[Error:] "<<mess<<endl;
		return 1;
	}

	GridBasedScoring* gbs = dynamic_cast<GridBasedScoring*>(scoring_function);
	for (list < Constraint* > ::iterator it = constraints.begin(); it != constraints.end(); it++)
	{
		scoring_function->constraints.push_back(*it);
		(*it)->setScoringFunction(scoring_function);
	}

	cout<<endl<<"-----------------------------------------"<<endl;
	cout<<"Scores will be calculated as : "<<scoring_function->getEquation()<<endl;
	cout<<"-----------------------------------------"<<endl<<endl;

	if (gbs != NULL)
	{
		gbs->replaceGridSetFromFile(grid_file);
	}

	Rescoring* rescoring = 0;

	if (!simple_rescoring)
	{
		if (!train && free_energy_label == "")
		{
			cerr<<"[Error:] free-energy label must be specified !"<<endl;
			return 1;
		}
		if (method == "Rescoring3D")
		{
			rescoring = new Rescoring3D(receptor, *ref_ligand, option, free_energy_label, scoring_function);
		}
		else if (method == "Rescoring4D")
		{
			rescoring = new Rescoring4D(receptor, *ref_ligand, option, free_energy_label, scoring_function);
		}
		else if (method == "Rescoring1D")
		{
			rescoring = new Rescoring1D(receptor, *ref_ligand, option, free_energy_label, scoring_function);
		}
		else
		{
			cerr<<"[Error:] Rescoring-method unknown !"<<endl;
			return 1;
		}
		if (!train)
		{
			rescoring->loadModel(model_file);
		}
		scoring_function->enableStoreInteractions();
	}

	scoring_function->setLigand(ref_ligand);
	scoring_function->update();
	scoring_function->updateScore();
	Log<<"====== Reference ligand ============"<<endl;
	scoring_function->printResult();


	/** If top fraction of docking results should not be rescored, then fetch scores and compute threshold for this fraction */

	bool ignore_top = false;
	double ignore_top_fraction = 0.0;
	if (par.has("tf"))
	{
		ignore_top_fraction = par.get("tf").toDouble();
		if (ignore_top_fraction < 1e-14 || ignore_top_fraction < 0 || ignore_top_fraction > 1)
		{
			ignore_top_fraction = 0.0;
			ignore_top = false;
		}
		else
		{
			ignore_top = true;
		}
	}
	double min_dock_score = -1e100;
	list<pair<double, bool> > rescore_list;
	if (ignore_top)
	{
		set<double> scores;
		GenericMolFile* input = MolFileFactory::open(par.get("i"));
		for (Molecule* mol = input->read(); mol; delete mol, mol = input->read())
		{
			if (mol->hasProperty("score"))
			{
				scores.insert(mol->getProperty("score").toString().toDouble());
			}
		}
		Size i = 0;
		Size max = scores.size()*ignore_top_fraction;
		set<double>::iterator s_it = scores.begin();
		for (; i < max; s_it++)
		{
			i++;
		}
		min_dock_score = *s_it;
		delete input;
	}


	/** Now, Rescore entire sd-/mol2-file   */

	double threshold = option.setDefaultReal("output_score_threshold", 1e100);

	GenericMolFile* input = MolFileFactory::open(par.get("i"));
	GenericMolFile* output = 0;

	if (simple_rescoring || !train)
	{
		output = MolFileFactory::open(par.get("o"), ios::out, input);

		DockResultFile* drf_output = dynamic_cast<DockResultFile*>(output);
		if (drf_output)
		{
			String dummy = "0";
			BALL::Docking::Result::Method method = Result::getMethod(3);
			String description = "";
			description = rescoring->getName()+"+"+scoring_function->getName();
			drf_output->setOutputParameters(method, "re-score", dummy, description);
		}
	}

	// Do the actual work ..
	processMultiMoleculeFile(scoring_function, sp, par_file, rescoring, train, min_dock_score, "score", input, output, threshold, ignore_top);

	if(!simple_rescoring)
	{
		if (train)
		{
			rescoring->recalibrate();
			rescoring->saveModel(par.get("o"));
		}
		else
		{
			double correlation, q2, stderr;
			rescoring->calculateQuality(correlation, q2, stderr);
			cout<<"Correlation = "<<String(correlation)<<endl;
			cout<<"Q2="<<String(q2)<<endl;
			cout<<"Standard error="<<String(stderr)<<endl;
		}
	}

	for (list<Constraint*>::iterator it = constraints.begin(); it != constraints.end(); it++)
	{
		delete *it;
	}

	delete rescoring;
	delete sp;
	delete ref_ligand;

	return 0;
}
