/*
 * A program that determines consensus structure from multiple RNA sequence alignment.
 *
 * (c) 2020 Mathews Lab, University of Rochester Medical Center.
 * Written by Abhinav Mittal
 */


#include "AlignmentFold.h"
#include "../src/structure.h"
#include "../src/rna_library.h"


//! constructor
AlignmentFold_Interface::AlignmentFold_Interface()
{
	calcType = "Consensus structure prediction";
	alphabet = DT_MSA;
	temperature = 310.15;
	bp_cutoff = 0.3;
	quiet = false;
	useBracketNotation = false;
	maxLoop = 40;
	disablecoax = true;
	quickfold = false;
	windowSize = -1;
	percent = 10;
	maxStructures = 20;
}

bool AlignmentFold_Interface::run()
{	
	// Create a variable that handles errors.
	int error = 0;
	char numstr[21];
	if (!quiet) cout << "Initializing nucleic acids..." << flush;
	RNA *strand = new RNA(seqFile.c_str(), FILE_MSA, alphabet.c_str());
	strand->GetStructure()->SetBasePairingCutoff(bp_cutoff);
	
	strand->GetStructure()->tmp_ct_hp = new structure;
	strand->GetStructure()->tmp_ct_hp->SetThermodynamicDataTable(strand->GetStructure()->GetThermodynamicDataTable());

	strand->GetStructure()->tmp_ct_i = new structure;
	strand->GetStructure()->tmp_ct_i->SetThermodynamicDataTable(strand->GetStructure()->GetThermodynamicDataTable());

	strand->GetStructure()->FillPairingMatrix();
	strand->GetStructure()->FillGapMatrix();
	ErrorChecker<RNA>* checker = new ErrorChecker<RNA>(strand);
	error = checker->isErrorStatus();
	if (error == 0) { if (!quiet) cout << "done." << endl; }

	/*
	 * Set the window size, based on the length of the sequence given as input.
	 * Only do this if window size hasn't been set on the command line.
	 * Use method GetSequenceLength to get the length of the sequence.
	 * The window sizes in relation to the length are hardcoded values.
	 */
	if (windowSize == -1 && error == 0) {
		int length = strand->GetSequenceLength();
		windowSize =
			(length > 1200) ? 20 :
			(length > 800) ? 15 :
			(length > 500) ? 11 :
			(length > 300) ? 7 :
			(length > 120) ? 5 :
			(length > 50) ? 3 :
			2;
	}

	if (error == 0) {
		// Show a message saying that the main calculation has started.
		if (!quiet) cout << "Running AlignmentFold ..." << flush;

		// Create the progress monitor.
		TProgressDialog* progress = new TProgressDialog();
		if (!quiet) strand->SetProgress(*progress);
		
		// Do the main calculation and check for errors.
		int mainCalcError = strand->FoldSingleStrand(percent, maxStructures, windowSize, saveFile.c_str(), maxLoop, quickfold, true, disablecoax);
		error = checker->isErrorStatus(mainCalcError);

		// Delete the progress monitor.
		strand->StopProgress();
		delete progress;

		// If no error occurred, print a message saying that the main calculation is done.
		if (error == 0) { if (!quiet) cout << "done." << endl; }
	}

	/*
		 * Write a CT output file using the WriteCt method.
		 * After writing is complete, use the error checker's isErrorStatus method to check for errors.
	*/
	if (error == 0) {

		// Show a message saying that the CT file is being written.
		if (!quiet) cout << "Writing output ct file..." << flush;

		// Write the CT file and check for errors.
		int writeError;

		if (useBracketNotation) {
			const DotBracketFormat bracketFmt = DBN_FMT_MULTI_TITLE; //see DotBracketFormat.h
			// write DotBracket file
			writeError = strand->WriteDotBracket(ctFile.c_str(), -1/*write ALL*/, bracketFmt);
		}
		else
			// write CT file
			writeError = strand->WriteCt(ctFile.c_str());

		error = checker->isErrorStatus(writeError);

		// If no errors occurred, show a CT file writing completion message.
		if (error == 0) { if (!quiet) cout << "done." << endl; }
	}

//	if (useBracketNotation) strand->WriteDotBracket(ctFile.c_str());
//	else strand->WriteCt(ctFile.c_str());

	for (int i = 0; i < strand->GetStructure()->GetNumberofSequences(); i++)
	{
		delete strand->GetStructure()->sequences_from_alignment[i].sequence_without_gaps;
	}
	delete strand->GetStructure()->tmp_ct_hp;
	delete strand->GetStructure()->tmp_ct_i;

	delete checker;
	delete strand;

	// Print confirmation of run finishing.
	if (error == 0) { if (!quiet) cout << calcType << " complete." << endl; }
	else { cerr << calcType << " complete with errors." << endl; }

	return true;
}

bool AlignmentFold_Interface::parse(int argc, char** argv) 
{	
	// Create the command line parser and build in its required parameters.
	ParseCommandLine* parser = new ParseCommandLine("AlignmentFold");
	parser->addParameterDescription("alignment file", "The name of a file containing multiple sequence alignment of homologous RNA sequences in FASTA format.");
	parser->addParameterDescription("CT file", "The name of a CT file to which output will be written. If the --bracket flag is present, output will be written as a dot-bracket file.\nIf the file name is a hyphen (-), the structure will be written to standard output (STDOUT) instead of a file.");

	// Add the Alphabet option.
//	vector<string> alphabetOptions;
//	alphabetOptions.push_back("-a");
//	alphabetOptions.push_back("--alphabet");
//	parser->addOptionFlagsWithParameters(alphabetOptions, "Specify the name of a folding alphabet and associated nearest neighbor parameters. The alphabet is the prefix for the thermodynamic parameter files, e.g. \"rna\" for RNA parameters or \"dna\" for DNA parameters or a custom extended/modified alphabet. The thermodynamic parameters need to reside in the at the location indicated by environment variable DATAPATH. The default is \"rna\" (i.e. use RNA parameters). This option overrides the --DNA flag.");
	
	// Add the base pairing cutoff
	vector<string> bpcutoffOptions;
	bpcutoffOptions.push_back("-b");
	bpcutoffOptions.push_back("-B");
	bpcutoffOptions.push_back("--bpcutoff");
	parser->addOptionFlagsWithParameters(bpcutoffOptions, "Specify the base pairing cutoff. Default is 0.3");

	// Add the quickfold option.
	vector<string> quickfoldOptions;
	quickfoldOptions.push_back("-mfe");
	quickfoldOptions.push_back("-MFE");
	quickfoldOptions.push_back("--MFE");
	parser->addOptionFlagsNoParameters(quickfoldOptions, "Specify that only the minimum free energy structure is needed.  No savefiles can be generated.  This saves nearly half the calculation time, but provides less information.");

	// Add the window size option.
	vector<string> windowOptions;
	windowOptions.push_back("-w");
	windowOptions.push_back("-W");
	windowOptions.push_back("--window");
	parser->addOptionFlagsWithParameters(windowOptions, "Specify a window size. Default is determined by the length of the sequence.");

	// Add the maximum loop size option.
	vector<string> loopOptions;
	loopOptions.push_back("-l");
	loopOptions.push_back("-L");
	loopOptions.push_back("--loop");
	parser->addOptionFlagsWithParameters(loopOptions, "Specify a maximum internal/bulge loop size. Default is 40 unpaired nucleotides.");

	// Add the percent energy difference option.
	vector<string> percentOptions;
	percentOptions.push_back("-p");
	percentOptions.push_back("-P");
	percentOptions.push_back("--percent");
	parser->addOptionFlagsWithParameters(percentOptions, "Specify a maximum percent energy difference. Default is 10 percent (specified as 10, not 0.1).");

	// Add the save file option.
	vector<string> saveOptions;
	saveOptions.push_back("-s");
	saveOptions.push_back("-S");
	saveOptions.push_back("--save");
	parser->addOptionFlagsWithParameters(saveOptions, "Specify the name of a save file, needed for dotplots and refolding. Default is not to generate a save file.");

	// Add the maximum number of structures option.
	vector<string> maxStructuresOptions;
	maxStructuresOptions.push_back("-m");
	maxStructuresOptions.push_back("-M");
	maxStructuresOptions.push_back("--maximum");
	parser->addOptionFlagsWithParameters(maxStructuresOptions, "Specify a maximum number of structures. Default is 20 structures.");

	// Add the disablecoax option.
	vector<string> disablecoaxOptions;
	disablecoaxOptions.push_back("--disablecoax");
	parser->addOptionFlagsNoParameters(disablecoaxOptions, "Specify that coaxial stacking recusions should not be used. This option uses a less realistic energy function in exchange for a faster calculation.");

	vector<string> quietOption = parser->addFlag(false, "-q --quite", "Suppress unnecessary output. This option is implied when the output file is '-' (STDOUT).");
	vector<string> bracketOption = parser->addFlag(false, "-k --bracket", "Write the predicted structure in dot-bracket notation (DBN) instead of CT format.");

	// Parse the command line into pieces.
	parser->parseLine(argc, argv);

	// Get required parameters from the parser.
	if (!parser->isError()) 
	{
		seqFile = parser->getParameter(1);
		ctFile = parser->getParameter(2);
	}

	quiet = parser->contains(quietOption) || isStdIoFile(ctFile.c_str()); // suppress unnecessary output if --quiet option is present or if the CT file is output to stdout.
	useBracketNotation = parser->contains(bracketOption); // write in dot-bracket notation.

	// Get the Alphabet option.
//	if (!parser->isError() && parser->contains(alphabetOptions))
//		alphabet = parser->getOptionString(alphabetOptions, false).c_str();

	
	// Get the bp_cutoff option.
	if (!parser->isError())
	{
		parser->setOptionFloat(bpcutoffOptions, bp_cutoff);
		if (bp_cutoff < .1)
		{
			parser->setError("bp_cutoff");
		}
	}

	// Get the maximum loop size option.
	if (!parser->isError()) {
		parser->setOptionInteger(loopOptions, maxLoop);
		if (maxLoop < 0) { parser->setError("maximum loop size"); }
	}

	// Get the quickfold option.
	if (!parser->isError()) { quickfold = parser->contains(quickfoldOptions); }

	// Get the window size option.
	if (!parser->isError()) {
		parser->setOptionInteger(windowOptions, windowSize);
		bool badWindow =
			(windowSize < 0) &&
			(windowSize != -1);
		if (badWindow) { parser->setError("window size"); }
	}

	// Get the maximum number of structures option.
	if (!parser->isError()) {
		parser->setOptionInteger(maxStructuresOptions, maxStructures);
		if (maxStructures <= 0) { parser->setError("maximum number of structures"); }
	}

	// Get the percent energy difference option.
	if (!parser->isError()) {
		parser->setOptionDouble(percentOptions, percent);
		if (percent < 0) { parser->setError("percent energy difference"); }
	}

	// Get the save file option.
	if (!parser->isError()) { saveFile = parser->getOptionString(saveOptions, false); }

	// Get the disablecoax option.
	// true by default, false if the flag is used
	if (!parser->isError()) { disablecoax = parser->contains(disablecoaxOptions); }

	// Delete the parser and return whether the parser encountered an error.
	bool noError = (parser->isError() == false);
	delete parser;
	return noError;
}

int main(int argc, char* argv[]) 
{
	AlignmentFold_Interface runner;
	if (!runner.parse(argc, argv)) return 1;
	return runner.run() ? 0 : 1;
}
