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

function BuildCodonFrequencies (obsF)
{
	PIStop = 1.0;
	model_dim = CountSenseCodons (_Genetic_Code);
	result = {model_dim,1};
	hshift = 0;

	for (h=0; h<64; 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];
		} else {
		    result[h-hshift][0]=obsF[first][0]*obsF[second][1]*obsF[third][2];
		}
	}
	return result*(1.0/PIStop);
}

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

global AC = 1;
global AT = 1;
global CG = 1;
global CT = 1;
global GT = 1;

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

_nucBiasTerms = {4,4};
_nucBiasTerms[0][0] = "";

hv = 0;

for (h=0; h<4; h=h+1)
{
	for (v=h+1; v<4; v=v+1)
	{
		_nucBiasTerms[h][v] = MGCustomRateBiasTerms[hv];
		_nucBiasTerms[v][h] = MGCustomRateBiasTerms[hv];
		hv = hv + 1;	
	}
}

h=0;
v=0;

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

function PopulateModelMatrix (ModelMatrixName&, EFV, synrateP, globalP, nonsynRateP) {

    ModelMatrixDimension = CountSenseCodons (_Genetic_Code);
	ModelMatrixName      = {ModelMatrixDimension,ModelMatrixDimension}; 

	hshift = 0;
	
	modelDefString = "";
	modelDefString*16384;
	
	catCounterAL = {};
	
	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];
				if (aa1==aa2) 
				{
					modelDefString*("ModelMatrixName["+hs+"]["+vs+"] := "+_nucBiasTerms[transition][transition2]+""+synrateP+"*EFV__["+ts+"]["+ps+"];\n"+
									"ModelMatrixName["+vs+"]["+hs+"] := "+_nucBiasTerms[transition][transition2]+""+synrateP+"*EFV__["+ts2+"]["+ps+"];\n");
				}
				else
				{
					if (Abs(globalP))
					{
						modelDefString*("ModelMatrixName["+hs+"]["+vs+"] := "+_nucBiasTerms[transition][transition2]+""+synrateP+"*"+globalP+"*EFV__["+ts+"]["+ps+"];\n"+
										"ModelMatrixName["+vs+"]["+hs+"] := "+_nucBiasTerms[transition][transition2]+""+synrateP+"*"+globalP+"*EFV__["+ts2+"]["+ps+"];\n");	
					
					}
					else
					{
						modelDefString*("ModelMatrixName["+hs+"]["+vs+"] := "+_nucBiasTerms[transition][transition2]+""+nonsynRateP+"*EFV__["+ts+"]["+ps+"];\n"+
										"ModelMatrixName["+vs+"]["+hs+"] := "+_nucBiasTerms[transition][transition2]+""+nonsynRateP+"*EFV__["+ts2+"]["+ps+"];\n");	
					}
				}
			}
	    }
    }		
	modelDefString*0;
	ExecuteCommands (modelDefString);
	return 0;
}

//--------------------------------------------------------------------------------------------------------------

lfunction extractBranchInformation (treeID, _omega_pattern, _model_prefix, _freq_prefix, _freq_vector, _do_mult) {
    _branchNames   = BranchName(^treeID,-1);
    _branchCount   = Columns (_branchNames) - 1;

    _toReturn         = {"Names": {}, "Props": {_branchCount,2}};  // cat count, branch length
    _brLenExpressions = {};
    _brLenLocals      = {};
    
    for (_i = 0; _i < _branchCount; _i += 1) {
        _expression = treeID + "\\." + _branchNames[_i] + "\\." + _omega_pattern;
        ExecuteCommands("GetInformation (_matchedVars, \"`_expression`\")");
        _catCount = Columns (_matchedVars);
        _toReturn ["Names"] + _branchNames[_i];
        (_toReturn["Props"])[_i][0] = _catCount;
        
        if (_catCount > 0) {
            _brLen     = 0;
            _totalWeight = 0;
            for (_c = 1; _c <= _catCount; _c += 1) {
                _modelID = _model_prefix + _c;
                if (Abs (_brLenExpressions[_modelID]) == 0) {
                    Model _internal = (^_modelID, ^_freq_vector, _do_mult);
                    GetString (_bl, _internal, -1);
                    _brLenExpressions [_modelID] = _bl;
                    _localParameters = {};
                    _p = 1;
                    GetString (_locP, _internal, 0);
                    while (_locP != None) {
                        _localParameters + _locP; 
                        GetString (_locP, _internal, _p);
                        _p += 1;
                    }
                    _brLenLocals [_modelID] = _localParameters;
                }
                                
                for (_lpc = 0; _lpc < Abs (_brLenLocals [_modelID]); _lpc += 1) {
                    local_var = (_brLenLocals [_modelID])[_lpc];
                    var_ref   =  treeID + "." + _branchNames[_i] + "." + (_brLenLocals [_modelID])[_lpc];
                    Eval (local_var + "=" + ^var_ref);
                }
                _thisModelLength = Eval (_brLenExpressions[_modelID]);
                
                if (_c < _catCount) {
                    var_ref = treeID + "." + _branchNames[_i] + "." + _freq_prefix + _c;
                    _freqW = ^var_ref;
                } else {
                    var_ref = treeID + "." + _branchNames[_i] + "." + _freq_prefix + (_c-1);
                    _freqW = 1.;
                }
                
                _thisModelWeight = (1-_totalWeight)*_freqW;            
                _totalWeight += _thisModelWeight;
                _brLen += _thisModelWeight * _thisModelLength;            
            }
        } else {
            _brLen = 3 * BranchLength (^treeID, _branchNames[_i]);
        }
        
        (_toReturn["Props"]) [_i][1] = _brLen/3;
    }
    
    
    return _toReturn;   
}

lfunction extractBranchLengthsFromBSREL (treeID) {
    _mleLengths  = extractBranchInformation (treeID, "omega[0-9]+", "MGMatrix", "Paux", "codon3x4", 0);
    _brLenByName = {};
     for (k = 0; k < Abs (_mleLengths["Names"]); k+=1) {
        _brLenByName[(_mleLengths["Names"])[k]] = (_mleLengths["Props"])[k][1];
    }
    return _brLenByName;
    
}

lfunction extractBranchLengthsFromBSREL_SRV (treeID) {
    _mleLengths  = extractBranchInformation (treeID, "nonsyn[0-9]+", "MGMatrix", "Paux", "codon3x4", 0);
    _brLenByName = {};
     for (k = 0; k < Abs (_mleLengths["Names"]); k+=1) {
        _brLenByName[(_mleLengths["Names"])[k]] = (_mleLengths["Props"])[k][1];
    }
    return _brLenByName;
    
}
