package dr.app.beagle.evomodel.substmodel;

import java.util.logging.Logger;

import dr.evolution.datatype.AminoAcids;
import dr.evolution.datatype.Codons;
import dr.evolution.datatype.GeneticCode;
import dr.evolution.datatype.Nucleotides;
import dr.app.beagle.evomodel.substmodel.BaseSubstitutionModel;
import dr.app.beagle.evomodel.substmodel.DefaultEigenSystem;
import dr.app.beagle.evomodel.substmodel.EigenSystem;
import dr.app.beagle.evomodel.substmodel.FrequencyModel;
import dr.app.beagle.evomodel.parsers.EmpiricalCodonModelParser;
import dr.inference.model.Parameter;

/**
 * Empirical model of codon evolution
 * 
 * @author Stefan Zoller
 */
public class EmpiricalCodonModel extends BaseSubstitutionModel {

	protected byte[] rateMap;
	
	protected Codons codonDataType;
	protected GeneticCode geneticCode;
	
	private Parameter omegaParameter;
	private Parameter kappaParameter;	// 2d: kappats and kappatv, 9d: ECM+F+omega
	private Parameter multintParameter;
	private EmpiricalRateMatrix rateMat;
	
	private int modelType;
	private final int ECM_OMEGA_2K = 2;
	private final int ECM_OMEGA_9K = 3;
	private final int ECM_OMEGA_NU = 4;
	private final int ECM_OMEGA = 1;


	/**
     * constructors
     *
     * @param codonDataType				Data type as Codons.UNIVERSAL
     * @param omegaParam				Parameter: Omega
     * @param kappaParam				Parameter: Kappa (multidimensional)
     * @param mntParam					Parameter: Multi-nt
     * @param rMat						Initial rate matrix and frequencies
     * @param freqModel					Frequency model
     */
	public EmpiricalCodonModel(Codons codonDataType,
		    Parameter omegaParam,
		    Parameter kappaParam,
		    Parameter mntParam,
		    EmpiricalRateMatrix rMat,
		    FrequencyModel freqModel) {
		this(codonDataType, omegaParam, kappaParam, mntParam, rMat, freqModel,
				new DefaultEigenSystem(codonDataType.getStateCount()));
	}
	
	
	public EmpiricalCodonModel(Codons codonDataType,
						    Parameter omegaParam,
						    Parameter kappaParam,
						    Parameter mntParam,
						    EmpiricalRateMatrix rMat,
						    FrequencyModel freqModel,
						    EigenSystem eigenSystem)
	{
		super(EmpiricalCodonModelParser.EMPIRICAL_CODON_MODEL, codonDataType, freqModel, eigenSystem);
		
		this.codonDataType = codonDataType;
		this.geneticCode = codonDataType.getGeneticCode();
		
		// setup parameters
		this.omegaParameter = omegaParam;
		addVariable(omegaParameter);
		omegaParameter.addBounds(new Parameter.DefaultBounds(Double.POSITIVE_INFINITY, 0.0,
				omegaParameter.getDimension()));
		
		if(kappaParam != null) {
			this.kappaParameter = kappaParam;
			addVariable(kappaParameter);
			kappaParameter.addBounds(new Parameter.DefaultBounds(Double.POSITIVE_INFINITY, 0.0,
				kappaParameter.getDimension()));
		}
		if(mntParam != null) {	
			this.multintParameter = mntParam;
			addVariable(multintParameter);
			multintParameter.addBounds(new Parameter.DefaultBounds(Double.POSITIVE_INFINITY, 0.0,
				multintParameter.getDimension()));
		}
		
		this.rateMat = rMat;

		constructRateMap();
		
		checkForModelType();
	}
	
	// decide which model to use: ECM_OMEGA_2K, ECM_OMEGA_9K, ECM_OMEGA_NU or ECM_OMEGA
	private void checkForModelType() {
		this.modelType = 0;
		if(kappaParameter != null) {
			if(kappaParameter.getDimension() == 2) {
				this.modelType = ECM_OMEGA_2K;
				Logger.getLogger("dr.evomodel").info("Using model ECM+omega+2k");
			} else {
				this.modelType = ECM_OMEGA_9K;
				Logger.getLogger("dr.evomodel").info("Using model ECM+omega+9k");
			}
		}
		if(multintParameter != null){
			this.modelType = ECM_OMEGA_NU;
			Logger.getLogger("dr.evomodel").info("Using model ECM+omega+nu");
		}
		if(kappaParameter == null && multintParameter == null) {
			this.modelType = ECM_OMEGA;
			Logger.getLogger("dr.evomodel").info("Using model ECM+omega");
		}
	}
    
	// setup substitution matrix depending on model type
	public void setupRelativeRates(double[] rates) {
		switch(modelType) {
		case ECM_OMEGA:
			setupRelativeRatesECMOmega(rates);
			break;
		case ECM_OMEGA_2K:
			setupRelativeRatesECMOmega2k(rates);
			break;
		case ECM_OMEGA_9K:
			setupRelativeRatesECMOmega9k(rates);
			break;
		case ECM_OMEGA_NU:
			setupRelativeRatesECMOmegaNu(rates);
			break;
		}
	}
	
	// actual setup routines for different models
	private void setupRelativeRatesECMOmega(double[] rates) {
		double[] initRateMatrix = rateMat.getRates();
		double omega = getOmega();
		for (int i = 0; i < rateCount; i++) {
			switch (rateMap[i]) {
			case 1:											// 1ts, 0tv, syn
			case 3:											// 0ts, 1tv, syn
			case 5:											// 2ts, 0tv, syn
			case 7:											// 1ts, 1tv, syn
			case 9:											// 0ts, 2tv, syn
			case 11:										// 3ts, 0tv, syn
			case 13:										// 2ts, 1tv, syn
			case 15:										// 1ts, 2tv, syn
			case 17: rates[i] = initRateMatrix[i]; break;	// 0ts, 3tv, syn
			
			case 2:													// 1ts, 0tv, nonsyn
			case 4:													// 0ts, 1tv, nonsyn
			case 6:													// 2ts, 0tv, nonsyn
			case 8:													// 1ts, 1tv, nonsyn
			case 10:												// 0ts, 2tv, nonsyn
			case 12:												// 3ts, 0tv, nonsyn
			case 14:												// 2ts, 1tv, nonsyn
			case 16:												// 1ts, 2tv, nonsyn
			case 18: rates[i] = initRateMatrix[i] * omega; break;	// 0ts, 3tv, nonsyn
			}
		}
	}
	
	private void setupRelativeRatesECMOmega2k(double[] rates) {
		double[] initRateMatrix = rateMat.getRates();
		double omega = getOmega();
		double kts = getKappaTs();
		double ktv = getKappaTv();
		for (int i = 0; i < rateCount; i++) {
			switch (rateMap[i]) {
				case 1: rates[i] = initRateMatrix[i] * kts; break;						// 1ts, 0tv, syn
				case 2: rates[i] = initRateMatrix[i] * kts * omega; break;				// 1ts, 0tv, nonsyn
				case 3: rates[i] = initRateMatrix[i] * ktv; break;						// 0ts, 1tv, syn
				case 4: rates[i] = initRateMatrix[i] * ktv * omega; break;				// 0ts, 1tv, nonsyn
				case 5: rates[i] = initRateMatrix[i] * kts * kts; break;				// 2ts, 0tv, syn
				case 6: rates[i] = initRateMatrix[i] * kts * kts * omega; break;		// 2ts, 0tv, nonsyn
				case 7: rates[i] = initRateMatrix[i] * kts * ktv; break;				// 1ts, 1tv, syn
				case 8: rates[i] = initRateMatrix[i] * kts * ktv * omega; break;		// 1ts, 1tv, nonsyn
				case 9: rates[i] = initRateMatrix[i] * ktv * ktv; break;				// 0ts, 2tv, syn
				case 10: rates[i] = initRateMatrix[i] * ktv * ktv * omega; break;		// 0ts, 2tv, nonsyn
				case 11: rates[i] = initRateMatrix[i] * kts * kts * kts; break;			// 3ts, 0tv, syn
				case 12: rates[i] = initRateMatrix[i] * kts * kts * kts * omega; break;	// 3ts, 0tv, nonsyn
				case 13: rates[i] = initRateMatrix[i] * kts * kts * ktv; break;			// 2ts, 1tv, syn
				case 14: rates[i] = initRateMatrix[i] * kts * kts * ktv * omega; break;	// 2ts, 1tv, nonsyn
				case 15: rates[i] = initRateMatrix[i] * kts * ktv * ktv; break;			// 1ts, 2tv, syn
				case 16: rates[i] = initRateMatrix[i] * kts * ktv * ktv * omega; break;	// 1ts, 2tv, nonsyn
				case 17: rates[i] = initRateMatrix[i] * ktv * ktv * ktv; break;			// 0ts, 3tv, syn
				case 18: rates[i] = initRateMatrix[i] * ktv * ktv * ktv * omega; break;	// 0ts, 3tv, nonsyn
			}
		}
	}
	
	private void setupRelativeRatesECMOmega9k(double[] rates) {
		double[] initRateMatrix = rateMat.getRates();
		double omega = getOmega();
		double[] kappa = getKappa();
		for (int i = 0; i < rateCount; i++) {
			switch (rateMap[i]) {
				case 1: rates[i] = initRateMatrix[i] * kappa[0]; break;						// 1ts, 0tv, syn
				case 2: rates[i] = initRateMatrix[i] * kappa[0] * omega; break;				// 1ts, 0tv, nonsyn
				case 3: rates[i] = initRateMatrix[i] * kappa[1]; break;						// 0ts, 1tv, syn
				case 4: rates[i] = initRateMatrix[i] * kappa[1] * omega; break;				// 0ts, 1tv, nonsyn
				case 5: rates[i] = initRateMatrix[i] * kappa[2]; break;				// 2ts, 0tv, syn
				case 6: rates[i] = initRateMatrix[i] * kappa[2] * omega; break;		// 2ts, 0tv, nonsyn
				case 7: rates[i] = initRateMatrix[i] * kappa[3]; break;				// 1ts, 1tv, syn
				case 8: rates[i] = initRateMatrix[i] * kappa[3] * omega; break;		// 1ts, 1tv, nonsyn
				case 9: rates[i] = initRateMatrix[i] * kappa[4]; break;				// 0ts, 2tv, syn
				case 10: rates[i] = initRateMatrix[i] * kappa[4] * omega; break;		// 0ts, 2tv, nonsyn
				case 11: rates[i] = initRateMatrix[i] * kappa[5]; break;			// 3ts, 0tv, syn
				case 12: rates[i] = initRateMatrix[i] * kappa[5] * omega; break;	// 3ts, 0tv, nonsyn
				case 13: rates[i] = initRateMatrix[i] * kappa[6]; break;			// 2ts, 1tv, syn
				case 14: rates[i] = initRateMatrix[i] * kappa[6] * omega; break;	// 2ts, 1tv, nonsyn
				case 15: rates[i] = initRateMatrix[i] * kappa[7]; break;			// 1ts, 2tv, syn
				case 16: rates[i] = initRateMatrix[i] * kappa[7] * omega; break;	// 1ts, 2tv, nonsyn
				case 17: rates[i] = initRateMatrix[i] * kappa[8]; break;			// 0ts, 3tv, syn
				case 18: rates[i] = initRateMatrix[i] * kappa[8] * omega; break;	// 0ts, 3tv, nonsyn
			}
		}
	}
	
	private void setupRelativeRatesECMOmegaNu(double[] rates) {
		double[] initRateMatrix = rateMat.getRates();
		double omega = getOmega();
		double mnt = getMultiNt();
		for (int i = 0; i < rateCount; i++) {
			switch (rateMap[i]) {
			case 1:														// 1ts, 0tv, syn
			case 3: rates[i] = initRateMatrix[i]; break;				// 0ts, 1tv, syn
			case 2: 													// 1ts, 0tv, nonsyn
			case 4: rates[i] = initRateMatrix[i] * omega; break;		// 0ts, 1tv, nonsyn
			
			case 5:														// 2ts, 0tv, syn
			case 7: 													// 1ts, 1tv, syn
			case 9: 													// 0ts, 2tv, syn
			case 11: 													// 3ts, 0tv, syn
			case 13: 													// 2ts, 1tv, syn
			case 15: 													// 1ts, 2tv, syn
			case 17: rates[i] = initRateMatrix[i] * mnt; break;			// 0ts, 3tv, syn
			
			case 6: 													// 2ts, 0tv, nonsyn
			case 8: 													// 1ts, 1tv, nonsyn
			case 10: 													// 0ts, 2tv, nonsyn
			case 12: 													// 3ts, 0tv, nonsyn
			case 14: 													// 2ts, 1tv, nonsyn
			case 16: 													// 1ts, 2tv, nonsyn
			case 18: rates[i] = initRateMatrix[i] * mnt * omega; break;	// 0ts, 3tv, nonsyn
			}
		}
	}
	
    protected void ratesChanged() {
	}
    
    protected void frequenciesChanged() {
	}
    
    // getter and setter for parameters

    public void setOmega(double omega) {
		omegaParameter.setParameterValue(0, omega);
		updateMatrix = true;
	}

    public double getOmega() { return omegaParameter.getParameterValue(0); }
    
    
    public void setKappa(double kts, double ktv) {
    	if(kappaParameter != null) {
    		kappaParameter.setParameterValue(0, kts);
    		kappaParameter.setParameterValue(1, ktv);
    		updateMatrix = true;
    	}
	}
    
    public double getKappaTs() { 
    	if(kappaParameter != null) {
    		return kappaParameter.getParameterValue(0);
    	} else {
    		return 0.0;
    	}
    }
    
    public double getKappaTv() { 
    	if(kappaParameter != null) {
    		return kappaParameter.getParameterValue(1);
    	} else {
    		return 0.0;
    	}
    }
    
    public double[] getKappa() {
    	if(kappaParameter != null) {
    		return kappaParameter.getParameterValues();
    	} else {
    		return new double[9];
    	}
    }
    
    
    public void setMultiNt(double mnt) {
    	if(multintParameter != null) {
    		multintParameter.setParameterValue(0, mnt);
    		updateMatrix = true;
    	}
	}

    public double getMultiNt() { 
    	if(multintParameter != null) {
    		return multintParameter.getParameterValue(0);
    	} else {
    		return 0.0;
    	}
	}
    
    
    /**
	 * Construct a map of the rate classes in the rate matrix using the current
	 * genetic code. Classes are:
	 * 	1-2: 1ts, 0tv (syn/nonsyn)
	 *  3-4: 0ts, 1tv
	 *  5-6: 2ts, 0tv
	 *  7-8: 1ts, 1tv
	 *  9-10: 0ts, 2tv
	 *  11-12: 3ts, 0tv
	 *  13-14: 2ts, 1tv
	 *  15-16: 1ts, 2tv
	 *  17-18: 0ts, 3tv
	 */
	protected void constructRateMap()
	{
		int u, v, i1, j1, k1, i2, j2, k2, ts, tv, non;
		byte rateClass;
		int[] codon;
		int cs1, cs2, aa1, aa2;

		int i = 0;

		rateMap = new byte[rateCount];

		for (u = 0; u < stateCount; u++) {

			codon = codonDataType.getTripletStates(u);
			i1 = codon[0];
			j1 = codon[1];
			k1 = codon[2];

			cs1 = codonDataType.getState(i1, j1, k1);
			aa1 = geneticCode.getAminoAcidState(codonDataType.getCanonicalState(cs1));

			for (v = u + 1; v < stateCount; v++) {
				
				ts = 0;
				tv = 0;
				non = 0;
				rateClass = -1;

				codon = codonDataType.getTripletStates(v);
				i2 = codon[0];
				j2 = codon[1];
				k2 = codon[2];

				cs2 = codonDataType.getState(i2, j2, k2);
				aa2 = geneticCode.getAminoAcidState(codonDataType.getCanonicalState(cs2));

				if (i1 != i2) {
					if ( (i1 == 0 && i2 == 2) || (i1 == 2 && i2 == 0) || // A <-> G
						 (i1 == 1 && i2 == 3) || (i1 == 3 && i2 == 1) ) { // C <-> T
						ts += 1; // Transition at position 1
					} else {
						tv += 1; // Transversion at position 1
					}
				}
				if (j1 != j2) {
					if ( (j1 == 0 && j2 == 2) || (j1 == 2 && j2 == 0) || // A <-> G
						(j1 == 1 && j2 == 3) || (j1 == 3 && j2 == 1) ) { // C <-> T
						ts += 1; // Transition
					} else {
						tv += 1; // Transversion
					}
				}
				if (k1 != k2) {
					if ( (k1 == 0 && k2 == 2) || (k1 == 2 && k2 == 0) || // A <-> G
						(k1 == 1 && k2 == 3) || (k1 == 3 && k2 == 1) ) { // C <-> T
						ts += 1; // Transition
					} else {
						tv += 1; // Transversion
					}
				}

	 			if (aa1 != aa2) {
					non = 1; // Is a non-synonymous change
				}

	 			// decide for rateClass
	 			switch(ts) {
	 				case 0:
	 					switch(tv) {
	 						case 1: rateClass = 3; break;	// 0ts, 1tv
							case 2: rateClass = 9; break;	// 0ts, 2tv
							case 3: rateClass = 17; break;	// 0ts, 3tv
							default: break;
	 					}
	 					break;
	 				case 1:
	 					switch(tv) {
	 						case 0: rateClass = 1; break;	// 1ts, 0tv
							case 1: rateClass = 7; break;	// 1ts, 1tv
							case 2: rateClass = 15; break;	// 1ts, 2tv
							default: break;
	 					}
	 					break;
	 				case 2:
	 					switch(tv) {
 							case 0: rateClass = 5; break;	// 2ts, 0tv
 							case 1: rateClass = 13; break;	// 2ts, 1tv
 							default: break;
	 					}
	 					break;
	 				case 3:
	 					rateClass = 11; break;	// 3ts, 0tv
	 				default: break;
	 			}
	 			
	 			if(non == 1) {
	 				rateClass += 1;
	 			}
				rateMap[i] = rateClass;
				i++;
			}

		}
	}

	public void printRateMap()
    {
        int u, v, i1, j1, k1, i2, j2, k2, ts, tv, non;
        byte rateClass;
        int[] codon;
        int cs1, cs2, aa1, aa2;

        System.out.print("\t");
        for (v = 0; v < stateCount; v++) {
            codon = codonDataType.getTripletStates(v);
            i2 = codon[0];
            j2 = codon[1];
            k2 = codon[2];

            System.out.print("\t" + Nucleotides.INSTANCE.getChar(i2));
            System.out.print(Nucleotides.INSTANCE.getChar(j2));
            System.out.print(Nucleotides.INSTANCE.getChar(k2));
        }
        System.out.println();

        System.out.print("\t");
        for (v = 0; v < stateCount; v++) {
            codon = codonDataType.getTripletStates(v);
            i2 = codon[0];
            j2 = codon[1];
            k2 = codon[2];

            cs2 = codonDataType.getState(i2, j2, k2);
            aa2 = geneticCode.getAminoAcidState(codonDataType.getCanonicalState(cs2));
            System.out.print("\t" + AminoAcids.INSTANCE.getChar(aa2));
        }
        System.out.println();

        for (u = 0; u < stateCount; u++) {

            codon = codonDataType.getTripletStates(u);
            i1 = codon[0];
            j1 = codon[1];
            k1 = codon[2];

            System.out.print(Nucleotides.INSTANCE.getChar(i1));
            System.out.print(Nucleotides.INSTANCE.getChar(j1));
            System.out.print(Nucleotides.INSTANCE.getChar(k1));

            cs1 = codonDataType.getState(i1, j1, k1);
            aa1 = geneticCode.getAminoAcidState(codonDataType.getCanonicalState(cs1));

            System.out.print("\t" + AminoAcids.INSTANCE.getChar(aa1));

            for (v = 0; v < stateCount; v++) {
	
            		ts = 0;
            		tv = 0;
            		non = 0;
	            	rateClass = -1;
	            	
	                codon = codonDataType.getTripletStates(v);
	                i2 = codon[0];
	                j2 = codon[1];
	                k2 = codon[2];
	
	                cs2 = codonDataType.getState(i2, j2, k2);
	                aa2 = geneticCode.getAminoAcidState(codonDataType.getCanonicalState(cs2));
	
	                if (i1 != i2) {
						if ( (i1 == 0 && i2 == 2) || (i1 == 2 && i2 == 0) || // A <-> G
							 (i1 == 1 && i2 == 3) || (i1 == 3 && i2 == 1) ) { // C <-> T
							ts += 1; // Transition at position 1
						} else {
							tv += 1; // Transversion at position 1
						}
					}
					if (j1 != j2) {
						if ( (j1 == 0 && j2 == 2) || (j1 == 2 && j2 == 0) || // A <-> G
							(j1 == 1 && j2 == 3) || (j1 == 3 && j2 == 1) ) { // C <-> T
							ts += 1; // Transition
						} else {
							tv += 1; // Transversion
						}
					}
					if (k1 != k2) {
						if ( (k1 == 0 && k2 == 2) || (k1 == 2 && k2 == 0) || // A <-> G
							(k1 == 1 && k2 == 3) || (k1 == 3 && k2 == 1) ) { // C <-> T
							ts += 1; // Transition
						} else {
							tv += 1; // Transversion
						}
					}
	
		 			if (aa1 != aa2) {
						non = 1; // Is a non-synonymous change
					}
	
		 			// decide for rateClass
		 			switch(ts) {
		 				case 0:
		 					switch(tv) {
		 						case 1: rateClass = 3; break;	// 0ts, 1tv
								case 2: rateClass = 9; break;	// 0ts, 2tv
								case 3: rateClass = 17; break;	// 0ts, 3tv
								default: break;
		 					}
		 					break;
		 				case 1:
		 					switch(tv) {
		 						case 0: rateClass = 1; break;	// 1ts, 0tv
								case 1: rateClass = 7; break;	// 1ts, 1tv
								case 2: rateClass = 15; break;	// 1ts, 2tv
								default: break;
		 					}
		 					break;
		 				case 2:
		 					switch(tv) {
	 							case 0: rateClass = 5; break;	// 2ts, 0tv
	 							case 1: rateClass = 13; break;	// 2ts, 1tv
	 							default: break;
		 					}
		 					break;
		 				case 3:
		 					rateClass = 11; break;	// 3ts, 0tv
		 				default: break;
		 			}
		 			
		 			if(non == 1) {
		 				rateClass += 1;
		 			}
	
	                System.out.print("\t" + rateClass);

            }
            System.out.println();

        }
    }
    // **************************************************************
    // XHTMLable IMPLEMENTATION
    // **************************************************************

	public String toXHTML() {
		StringBuffer buffer = new StringBuffer();

		buffer.append("<em>Empirical Codon Model</em> omega = ");
		buffer.append(getOmega());
		buffer.append(", kappa_ts = ");
		buffer.append(getKappaTs());
		buffer.append(", kappa_tv = ");
		buffer.append(getKappaTv());
		buffer.append(", multi_nt = ");
		buffer.append(getMultiNt());
		buffer.append(", initial matrix = " + rateMat.getDirName() + "/" + rateMat.getMatName());
		buffer.append(", initial freqs = " + rateMat.getDirName() + "/" + rateMat.getFreqName());
		return buffer.toString();
	}

	
	static String format = "%2.4e";
	
	public String printRelRates() {
		StringBuffer sb = new StringBuffer();
		for (int i = 0; i < relativeRates.length; i++) {
			sb.append(String.format(format, relativeRates[i]));
			sb.append("\t");
		}
		sb.append("\n");
		return sb.toString();
	}
}
