/*	this file defines the ECM model as reported in Kosiol, Holmes & Goldman (2007). Molecular Biology and Evolution 24:1464-1479 

	09/03/2009
	
	Coded by Gordon Botha
	Tweaked by Wayne Delport
*/

NUC_LETTERS = "ACGT";
AA = "FLIMVSPTAY*HQNKDECWRG";

CODONS = {64, 1};	
for (i = 0; i < 64; i = i + 1)
{
	CODONS [i] = NUC_LETTERS[i$16] + NUC_LETTERS [(i%16)$4] + NUC_LETTERS [i%4];
}

ECMmodelType = 1; /* 0 for ECM, 1 for ECM+F, 2 for ECM+F+w */

ModelMatrixDimension = 0;

function BuildCodonFrequencies (obsF)
{
	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][0]*obsF[third][0];
			continue; 
		}
		result[h-hshift][0]=obsF[first][0]*obsF[second][0]*obsF[third][0];
	}
	return result*(1.0/PIStop);
}

function PopulateModelMatrix (ModelMatrixName& )
{		

	ModelMatrixName = {ModelMatrixDimension,ModelMatrixDimension}; 
	hshift = 0;
	if (ECMmodelType > 1)		/* model where omega is implemented for the first time, ECM + F + omega */
	{
		global globalVariable_omega = 0.192;		/* value initialised as omega_E, p 1476 of paper */
	}
	else {
		global globalVariable_omega := 1;
	}
	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)
		{
			if (_Genetic_Code[v]==10) 
			{
				vshift = vshift+1;
				continue; 
			}
			hs = h-hshift;
			vs = v-vshift;
			if ( _Genetic_Code [ v ] == _Genetic_Code [ h ] )
			{
				ModelMatrixName[hs][vs] := _numericRateMatrix__[hs__][vs__]*t;
				ModelMatrixName[vs][hs] := _numericRateMatrix__[hs__][vs__]*t;
			}
			else
			{
				ModelMatrixName[hs][vs] := globalVariable_omega*_numericRateMatrix__[hs__][vs__]*t;
				ModelMatrixName[vs][hs] := globalVariable_omega*_numericRateMatrix__[hs__][vs__]*t;
			}
		}
	}
	return 1;
}

fprintf ( stdout, "\n\nAllow multiple substitutions between codons? (Y/N)" );
fscanf ( stdin, "String", msubs );

if ( ( msubs == "Y" ) || ( msubs == "y" )  ) {
	incFileName = HYPHY_LIB_DIRECTORY + "TemplateBatchFiles" + DIRECTORY_SEPARATOR + "TemplateModels" + DIRECTORY_SEPARATOR + "EmpiricalCodon" + DIRECTORY_SEPARATOR + "KHG_ECMu";
}
else {
	incFileName = HYPHY_LIB_DIRECTORY + "TemplateBatchFiles" + DIRECTORY_SEPARATOR + "TemplateModels" + DIRECTORY_SEPARATOR + "EmpiricalCodon" + DIRECTORY_SEPARATOR + "KHG_ECM";
}

fscanf ( incFileName, "Matrix,NMatrix,NMatrix", KHG_codonOrder, KHG_RateMatrix, KHG_vof );

ModelMatrixDimension = 0;
for ( i = 0; i < 64; i = i + 1 ) {
	if ( _Genetic_Code [ i ] != 10 ) {
		ModelMatrixDimension = ModelMatrixDimension + 1;
	}
}

/*array of codon indeces in the KHG model*/
CodonCompare = { ModelMatrixDimension, 1 };
ishit = 0;
for ( i = 0; i < 64; i = i + 1 ) {
	if ( _Genetic_Code [ i ] == 10 ) {
		ishit = ishit + 1;
	}	
	else {
		query = CODONS [ i ];
		got = 0;
		j = 0;
		while ( ( got == 0 ) && ( j < Columns(KHG_codonOrder) ) ) {
			if ( query == KHG_codonOrder [ j ] ) {
				/* fprintf ( stdout, query, " ", KHG_codonOrder [ j ], "\n" ); */
				got = 1;
				CodonCompare [ i-ishit ] = j;
			}
			j = j + 1;
		}
	}
}

_numericRateMatrix = { ModelMatrixDimension, ModelMatrixDimension }; /* codons to be ordered as in HYPHY*/
ECM_Freqs = { ModelMatrixDimension, 1 };
for ( i = 0; i < ModelMatrixDimension; i = i + 1 ) {
	ECM_Freqs [ i ] = KHG_vof [ CodonCompare [ i ] ];	
	for ( j = i + 1; j < ModelMatrixDimension; j = j + 1 ) {
		_numericRateMatrix [ i ][ j ] = KHG_RateMatrix [ CodonCompare [ i ] ][ CodonCompare [ j ] ];
		_numericRateMatrix [ j ][ i ] = KHG_RateMatrix [ CodonCompare [ j ] ][ CodonCompare [ i ] ];
	}
}

if ( ECMmodelType > 0 ) { /* Use F61 frequencies as per KHG */
	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! */
}
else {
	vectorOfFrequencies = ECM_Freqs;
}

ECM = 0;
PopulateModelMatrix ( "ECM" );
Model ECMModel = ( ECM, vectorOfFrequencies, 1 );


