/* this file defines the MEC model from 

A Combined Empirical and Mechanistic Codon Model
By Adi Doron-Faigenboim and Tal Pupko 
Mol. Biol. Evol. 24(2):388Ð397. 2007

v0.1 20100520 SLKP spond@ucsd.edu 

*/



global T_R   = 1; 	/* the global rate for 1 transition substitutions */
global T_V   = 1; 	/* the global rate for 1 transversion substitutions */
global T_RR  = 1;	/* the global rate for 2 transition substitutions */
global T_VV  = 1;	/* the global rate for 2 transversion substitutions */
global T_RV  = 1;	/* the global rate for 1 transversion/1 transversion substitutions */
global T_SUB = 1;	/* the global rate for 3-nucleotide substitutions */

global omega = 1; 	/* the global dN/dS rate */
global F	 = 0.5; /* the mixing proportion of the neutral/empirical model */

F:<1.0;

/*---------------------------------------------------------------------------------------------------------------------------------------------*/

function PopulateModelMatrix (ModelMatrixName&, CodonFreqs, AAFreqs, AAOrdering, AARates, categoriesUsed)

/* input
		ModelMatrixName : the ID to store the populated matrix into 
		CodonFreqs	    : Frequencies of sense codons
		
		AAFreqs			: Frequencies of the 20 amino-acids (in the _hyphy ordering including the stop codon; 
						  this is auto-generated by BuildCodonFrequencies )
						  
		AAOrdering		: The ordering of amino-acids in the rate matrix
		AARates			: A 20x20 rate matrix for the amino-acids
		categoriesUsed		: set to != 0 to include rate variation in omega
		
*/

{
	
	ModelMatrixName = {ModelMatrixDimension,ModelMatrixDimension}; 
	
	ModelMatrixName [0][0] = ""; /* a hack to force ModelMatrixName into strings */
		
	/* define a map from the internal HyPhy ordering of AA to what is found in AAFreqs and AARates*/
	
	_aaMapping = {};
	
	for (h = 0; h < 20; h = h+1)
	{
		_aaMapping [AAOrdering[h]] = h;
	}
	
	modelDefString = "";
	modelDefString * 256;
	
	substitutionTypeStrings = {{"T_R","T_V","T_RR","T_VV","T_RV","T_SUB"}};
	
	byAA	= {}; 
		/* store the contributions of each substitution type, index by AA pair */
	
	hshift = 0;
	for (h=0; h<64; h=h+1)
	{
		if (IsStop(h,_Genetic_Code)) 
		{
			hshift = hshift+1;
			continue; 
		}
		codeH = _hyphyAAOrdering[_Genetic_Code[h]];
		vshift = hshift;
		for (v = h+1; v<64; v=v+1)
		{
			if (IsStop(v,_Genetic_Code)) 
			{
				vshift = vshift+1;
				continue; 
			}
			
			codeV = _hyphyAAOrdering[_Genetic_Code[v]];
			
			codonProps = CompareCodonProperties (h,v,_Genetic_Code);
			
			tiCount = 0;
			tvCount = 0;
			
			for (_p = 1; _p <= 3; _p = _p+1)
			{
				_titv = IsTransition (codonProps[_p]);
				tiCount = tiCount + (_titv > 0);
				tvCount = tvCount + (_titv < 0);
			}
			
			if (tiCount + tvCount == 3)
			{
				idx = 5;
			}
			else
			{
				if (tiCount == 2)
				{
					idx = 2;
				}
				else
				{
					if (tvCount == 2)
					{
						idx = 3;
					}
					else
					{
						if (tiCount + tvCount == 2)
						{
							idx = 4;
						}
						else
						{
							if (tiCount == 1)
							{
								idx = 0;
							}
							else
							{
								idx = 1;
							}
						}
					}
				}
			}
			
			if (codeH != codeV)
			{
				if (codeH < codeV)
				{
					codeToken = codeH + codeV;
				}
				else
				{
					codeToken = codeV + codeH;
				}
				if (Rows(byAA[codeToken]) == 0) /* not seen this pair yet */
				{
					byAA[codeToken] = {1,6}; 
						/* 1-transition, 1-transversion, 2-transitions, 2-transversions,1-each,3-subs;
						   multiplicative factors in front of the global rates
						*/
				}	
				(byAA [codeToken])[idx] = (byAA [codeToken])[idx] + 2*CodonFreqs[h-hshift]*CodonFreqs[v-vshift];
			}
			ModelMatrixName[h-hshift][v-vshift] = "t*" + substitutionTypeStrings[idx];
	    }
    }	
    
    /* pass 2 to define amino-acid correction weights on a per-residue pair basis */
    aaExpressions = {};
    hshift = 0;
	for (h=0; h<64; h=h+1)
	{
		if (IsStop(h,_Genetic_Code)) 
		{
			hshift = hshift+1;
			continue; 
		}
		codeH = _hyphyAAOrdering[_Genetic_Code[h]];
		vshift = hshift;
		for (v = h+1; v<64; v=v+1)
		{
			if (IsStop(v,_Genetic_Code)) 
			{
				vshift = vshift+1;
				continue; 
			}
			
			codeV = _hyphyAAOrdering[_Genetic_Code[v]];  
			if (codeH != codeV)
			{
				if (codeH < codeV)
				{
					codeToken = codeH + codeV;
				}
				else
				{
					codeToken = codeV + codeH;
				}
				if (Abs(aaExpressions[codeToken]) == 0) /* not seen this pair yet */
				{
					exprStr = "";
					for (idx = 0; idx < 6; idx = idx+1)
					{
						if ((byAA[codeToken])[idx] > 0.)
						{
							exprStr = exprStr + substitutionTypeStrings[idx] + "*" + (byAA[codeToken])[idx] + "+"; 
						}
					}
					
					aaRate = AARates[_aaMapping[codeH]][_aaMapping[codeV]];
					aaExpressions[codeToken] = "" + aaRate * AAFreqs[_Genetic_Code[h]] + "/(" + exprStr[0][Abs(exprStr)-2] + ")";
				}	
				
				if (categoriesUsed)
				{
					ModelMatrixName[h-hshift][v-vshift] = ModelMatrixName[h-hshift][v-vshift] + "*(c*F*"+aaExpressions[codeToken]+"+(1-F))";
				}
				else
				{
					ModelMatrixName[h-hshift][v-vshift] = ModelMatrixName[h-hshift][v-vshift] + "*(omega*F*"+aaExpressions[codeToken]+"+(1-F))";
				}
			}
		}
	}
    
    modelDefString * "ModelMatrixName = {ModelMatrixDimension,ModelMatrixDimension}; "; 
    
    hshift = 0;
	for (h=0; h<64; h=h+1)
	{
		if (IsStop(h,_Genetic_Code)) 
		{
			hshift = hshift+1;
			continue; 
		}
		codeH = _hyphyAAOrdering[_Genetic_Code[h]];
		vshift = hshift;
		for (v = h+1; v<64; v=v+1)
		{
			if (IsStop(v,_Genetic_Code)) 
			{
				vshift = vshift+1;
				continue; 
			}
			
			modelDefString * ("ModelMatrixName["+(h-hshift)+"]["+(v-vshift)+"] := " + ModelMatrixName[h-hshift][v-vshift] + ";\n");
			modelDefString * ("ModelMatrixName["+(v-vshift)+"]["+(h-hshift)+"] := " + ModelMatrixName[h-hshift][v-vshift] + ";\n");
		}
	}
	
    modelDefString * 0;    
	ExecuteCommands (modelDefString);
   	return 1;

}

/*---------------------------------------------------------------------------------------------------------------------------------------------*/

function BuildCodonFrequencies (obsF, AAFreqs&)
{
	PIStop  = 1.0;
	result  = {ModelMatrixDimension,1};
	hshift  = 0;

	AAFreqs = {21,1};

	for (h=0; h<64; h=h+1)
	{
		first  = h$16;
		second = h%16$4;
		third  = h%4;
		myCode = _Genetic_Code[h];
		if (myCode==10) 
		{
			hshift = hshift+1;
			PIStop = PIStop-obsF[first][0]*obsF[second][1]*obsF[third][2];
			continue; 
		}
		result[h-hshift]=obsF[first][0]*obsF[second][1]*obsF[third][2];
		AAFreqs[myCode] = AAFreqs[myCode] + result[h-hshift];
	}
	
	AAFreqs = AAFreqs*(1.0/PIStop);
	return result*(1.0/PIStop);
}

/*---------------------------------------------------------------------------------------------------------------------------------------------*/

if (!SKIP_MODEL_PARAMETER_LIST)
{
	categoriesUsed = 0;
	ExecuteAFile ("TemplateModels/EmpiricalAA/modellist.ibf");
	
	models = generateSelectionList();
	
	ChoiceList (whichModel,"Protein Model Component", 1, SKIP_NONE,models );
	
	if (whichModel < 0)
	{
		return 0;
	}
	models = "TemplateModels/EmpiricalAA/"+ (modelList[whichModel])["File"];
	fscanf (models, "String,NMatrix", modelOrdering, modelMatrix);

	#include "modelParameters4.mdl";

	if (rateType >= 1)
	{
		ExecuteAFile ("TemplateModels/EmpiricalAA/modellist.ibf");
		ExecuteAFile ("Utility/GrabBag.bf");
		ExecuteAFile ("TemplateModels/Yang2000Distributions.def");
		rateType = rateType - 1;
		SetWDistribution (prompt_for_a_value ("\nNumber of rate categories",10,2,15,1));
		categoriesUsed = 1;
	}	
}


ModelMatrixDimension = 64 - (_Genetic_Code["_MATRIX_ELEMENT_VALUE_ == 10"]*Transpose(_Genetic_Code["1"]))[0];

if (!SKIP_HARVEST_FREQ)
{
	HarvestFrequencies (observedFreq,filteredData,3,1,1);
}

vectorOfFrequencies			= BuildCodonFrequencies (observedFreq, "AA");
MULTIPLY_BY_FREQS			= PopulateModelMatrix ("MECMatrix", vectorOfFrequencies, AA, modelOrdering, modelMatrix, categoriesUsed);
Model MEC_Model 			= (MECMatrix,vectorOfFrequencies,1);

USE_POSITION_SPECIFIC_FREQS = 1;
