/*	Implementation of the Linear Combination of Amino Acid Properties model by Conant, Wagner & Stadler 2006

	Conant, G.C. and P.F. Stadler, Solvent exposure imparts similar selective pressures across a range of yeast proteins. Mol Biol Evol, 2009. 26(5): p. 1155-61.
	Conant, G.C., G.P. Wagner, and P.F. Stadler, Modeling amino acid substitution patterns in orthologous and paralogous genes. Mol Phylogenet Evol, 2007. 42(2): p. 298-307.

	Thanks to Gavin Conant for supplying amino acid property matrices
	
	09/04/2009 by Wayne "die boer" Delport

*/

USE_ADAPTIVE_VARIABLE_STEP = 1;

numProps = 5; /* number of amino acid property matrices included in the LCAP file */

incFileName = HYPHY_LIB_DIRECTORY + "TemplateBatchFiles" + DIRECTORY_SEPARATOR + "TemplateModels" + DIRECTORY_SEPARATOR + "EmpiricalAA" + DIRECTORY_SEPARATOR + "LCAP";
fscanf ( incFileName, "NMatrix,NMatrix,NMatrix,NMatrix,NMatrix", prop_0, prop_1, prop_2, prop_3, prop_4 );

/* order as in LCAP file from Gavin Conant : Chemical Composition, Polarity, Volume, Iso-electric point, Hydropathy */

function BuildCodonFrequencies1x4 (obsF)
{
	if (!ModelMatrixDimension)
	{
		ModelMatrixDimension = 64 - (_Genetic_Code["_MATRIX_ELEMENT_VALUE_ == 10"]*Transpose(_Genetic_Code["1"]))[0];
	}

	PIStop = 1.0;
	result = {ModelMatrixDimension,1};
	hshift = 0;

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

function BuildCodonFrequencies3x4 (obsF)
{
	if (!ModelMatrixDimension)
	{
		ModelMatrixDimension = 64 - (_Genetic_Code["_MATRIX_ELEMENT_VALUE_ == 10"]*Transpose(_Genetic_Code["1"]))[0];
	}

	PIStop = 1.0;
	result = {ModelMatrixDimension,1};
	hshift = 0;

	for (h=0; h<64; h=h+1)
	{
		first = h$16;
		second = h%16$4;
		third = h%4;
		if (_Genetic_Code[h]==10) 
		{
			hshift = hshift+1;
			PIStop = PIStop-obsF[first][0]*obsF[second][1]*obsF[third][2];
			continue; 
		}
		result[h-hshift][0]=obsF[first][0]*obsF[second][1]*obsF[third][2];
	}
	return result*(1.0/PIStop);
}
												  
function PopulateModelMatrix (ModelMatrixName&, EFV, freqs)
{
	if (!ModelMatrixDimension)
	{
		ModelMatrixDimension = 64 - (_Genetic_Code["_MATRIX_ELEMENT_VALUE_ == 10"]*Transpose(_Genetic_Code["1"]))[0];
	}
		
	ModelMatrixName = {ModelMatrixDimension,ModelMatrixDimension}; 
	
	synCatRateMult 	  = "synRate*";
	nonsynCatRateMult = "synRate*";
	global C = 0.5; C :> -10000;
	
	modelDefString = "";
	modelDefString*16384;
	
	for ( aaprop = 0; aaprop < numProps; aaprop = aaprop + 1 ) {
		ExecuteCommands ( "global alpha_" + aaprop + " = 0.5; global alpha_" + aaprop + " :> -10000;" );
	}
	
	hshift = 0;
	
	for (h=0; h<64; h=h+1)
	{
		if (_Genetic_Code[h]==10) 
		{
			hshift = hshift+1;
			continue; 
		}
		vshift = hshift;
		for (v = h+1; v<64; v=v+1)
		{
			diff = v-h;
			if (_Genetic_Code[v]==10) 
			{
				vshift = vshift+1;
				continue; 
			}
			nucPosInCodon = 2;
			if ((h$4==v$4)||((diff%4==0)&&(h$16==v$16))||(diff%16==0))
			{
				if (h$4==v$4)
				{
					transition = v%4;
					transition2= h%4;
				}
				else
				{
					if(diff%16==0)
					{
						transition = v$16;
						transition2= h$16;
						nucPosInCodon = 0;
					}
					else
					{
						transition = v%16$4;
						transition2= h%16$4;
						nucPosInCodon = 1;
					}
				}
				hs = Format(h-hshift,0,0);
				vs = Format(v-vshift,0,0);
				ts = Format(transition,0,0);
				ts2= Format(transition2,0,0);
				ps = Format(nucPosInCodon,0,0);
				aa1 = _Genetic_Code[0][h];
				aa2 = _Genetic_Code[0][v];
				
				
				/*Conant does not have stop AA in the matrices, so we need to adjust for aa > 10 */
				if ( aa1 < 10 ) {
					caa1 = aa1;
				}
				else {
					caa1 = aa1 - 1;
				}
				if ( aa2 < 10 ) {
					caa2 = aa2;
				}
				else {
					caa2 = aa2 - 1;
				}
				
				/* add in the amino acid property multipliers */
				aapropMultiplier = "";
				for ( aaprop = 0; aaprop < numProps; aaprop = aaprop + 1 ) {
					propMultiString = "";
					propMultiString * ( "propMulti = Abs(prop_" + aaprop + "[" + caa1 + "][" + caa2 + "]);" );
					propMultiString*0;
					ExecuteCommands ( propMultiString );
					if ( aaprop == 0 ) {
						aapropMultiplier * ("Exp( C + alpha_" + aaprop + "*" + propMulti + "" );
					}
					else {
						aapropMultiplier * ("+ alpha_" + aaprop + "*" + propMulti + "" );
					}
					
				}
				aapropMultiplier * ( ")" );
				
				if ( freqs < 2 ) {
					if ( freqs == 0 ) {
						ps = 0;
					}
				
					if (aa1==aa2) 
					{
						modelDefString*("ModelMatrixName["+hs+"]["+vs+"] := "+synCatRateMult+_nucBiasTerms[transition][transition2]+"EFV__["+ts+"]["+ps+"];\n"+
													 "ModelMatrixName["+vs+"]["+hs+"] := "+synCatRateMult+_nucBiasTerms[transition][transition2]+"EFV__["+ts2+"]["+ps+"];\n");
					}
					else
					{
						modelDefString*("ModelMatrixName["+hs+"]["+vs+"] := "+nonsynCatRateMult+_nucBiasTerms[transition][transition2]+"EFV__["+ts+"]["+ps+"]*" + aapropMultiplier + ";\n"+
													 "ModelMatrixName["+vs+"]["+hs+"] := "+nonsynCatRateMult+_nucBiasTerms[transition][transition2]+"EFV__["+ts2+"]["+ps+"]*" + aapropMultiplier + ";\n");						
					}
				}
				else {
					if (aa1==aa2) 
					{
						modelDefString*("ModelMatrixName["+hs+"]["+vs+"] := "+synCatRateMult+_nucBiasTerms[transition][transition2]+"1;\n"+
													 "ModelMatrixName["+vs+"]["+hs+"] := "+synCatRateMult+_nucBiasTerms[transition][transition2]+"1;\n");
					}
					else
					{
						modelDefString*("ModelMatrixName["+hs+"]["+vs+"] := "+nonsynCatRateMult+_nucBiasTerms[transition][transition2]+aapropMultiplier + "*1;\n"+
													 "ModelMatrixName["+vs+"]["+hs+"] := "+nonsynCatRateMult+_nucBiasTerms[transition][transition2]+aapropMultiplier + "*1;\n");						
					}
				}
			}
	    }
    }		
	modelDefString*0;
	/*fprintf ( stdout, modelDefString );*/
	ExecuteCommands (modelDefString);
	
	return freqs == 2;
}

 ChoiceList (freqType,"Choose a frequency parameterisation",1,SKIP_NONE,
			"F1x4","Empirical nucleotide frequencies",
			"F3x4","Empirical position specific nucleotide frequencies",
			"F61","Empirical codon frequencies",
			"ML3x4","Maximum likelihood estimated position specific frequencies");


fprintf						  (stdout, "\nModel string for nucleotide biases: (eg: 010010 for HKY85)");
fscanf						  (stdin,"String",modelDesc);


MGCustomRateBiasTerms = {{"AC*","","AT*","CG*","CT*","GT*"}};	

		
paramCount	   = 0;
_nucBiasTerms  = {4,4};

_nucBiasTerms[0][0] = "";


if (modelDesc[0]==modelDesc[1])
{
	MGCustomRateBiasTerms[0] = MGCustomRateBiasTerms[1];
}

_nucBiasTerms[1][0] = MGCustomRateBiasTerms[0];
_nucBiasTerms[0][1] = MGCustomRateBiasTerms[0];
_nucBiasTerms[2][0] = MGCustomRateBiasTerms[1];
_nucBiasTerms[0][2] = MGCustomRateBiasTerms[1];

h = 0;
v = 3;

for (customLoopCounter2=2; customLoopCounter2<6; customLoopCounter2=customLoopCounter2+1)
{
	for (customLoopCounter=0; customLoopCounter<customLoopCounter2; customLoopCounter=customLoopCounter+1)
	{
		if (modelDesc[customLoopCounter]==modelDesc[customLoopCounter2])
		{
			_nucBiasTerms[h][v] = MGCustomRateBiasTerms[customLoopCounter];
			_nucBiasTerms[v][h] = MGCustomRateBiasTerms[customLoopCounter];
			break;
		}
	}
	if (customLoopCounter == customLoopCounter2)
	{
		_nucBiasTerms[h][v] = MGCustomRateBiasTerms[customLoopCounter2];
		_nucBiasTerms[v][h] = MGCustomRateBiasTerms[customLoopCounter2];
	}
	
	v = v+1;
	if (v==4)
	{
		h=h+1;
		v=h+1;
	}
}

prev = "";
for ( i = 0; i < Columns ( _nucBiasTerms ); i = i + 1 ) {
	for ( j = 0; j < Rows ( _nucBiasTerms ); j = j + 1 ) {
		if ( Abs ( _nucBiasTerms [i][j] ) ) {
			nucPair = _nucBiasTerms [i][j];
			if ( prev != nucPair ) {
				ExecuteCommands ( "global " + nucPair[0][Abs(nucPair)-2] + " = 1;" );
				prev = nucPair;
			}
		}
	}
}

if ( freqType < 1 ) {
	HarvestFrequencies (observedFreq,filteredData,1,1,1);
	vectorOfFrequencies = BuildCodonFrequencies1x4 (observedFreq);
}
else {
	if ( freqType == 1 ) { /* F3x4*/
		HarvestFrequencies ( observedFreq, filteredData, 3,1,1);
		vectorOfFrequencies = BuildCodonFrequencies3x4 (observedFreq);
	}
	if ( freqType == 2 ) { /* F61 */
 		COUNT_GAPS_IN_FREQUENCIES = 0;
		HarvestFrequencies (observedFreq,filteredData,3,3,1);
		vectorOfFrequencies = Transpose(observedFreq[Transpose(_Genetic_Code["_MATRIX_ELEMENT_VALUE_!=10"])]); /* cough, cough, cough ... HACK! */
	}
	if ( freqType == 3 ) { /*MLF3x4*/
		if (!ModelMatrixDimension)
		{
			ModelMatrixDimension = 64 - (_Genetic_Code["_MATRIX_ELEMENT_VALUE_ == 10"]*Transpose(_Genetic_Code["1"]))[0];
		}
		HarvestFrequencies ( observedFreq, filteredData, 3,1,1);
		#include "MGFreqsEstimator.ibf";
		BuildCodonFrequencies (paramFreqs,"vectorOfFrequencies");
	}
}

MULTIPLY_BY_FREQS = PopulateModelMatrix ("LCAP", observedFreq, freqType);
Model LCAPModel = (LCAP,vectorOfFrequencies,MULTIPLY_BY_FREQS);
Export ( modelString, LCAPModel );
fprintf ( stdout, modelString, "\n" );

fprintf ( stdout, "alpha parameters in order (0-4) are Chemical Composition, Polarity, Volume, Iso-electric point, Hydropathy\n" );
