package dr.app.beagle.evomodel.substmodel;

import dr.app.beagle.evomodel.sitemodel.SiteRateModel;
import dr.inference.model.Parameter;
import dr.math.KroneckerOperation;
import dr.app.beagle.evomodel.substmodel.*;


import java.util.List;

/**
 * A derived class from ProductChainSubstitutionModel that introduces dependence between sites.
 *
 * @version 08/01/2010
 *
 * @author Marc A. Suchard
 * @author Yu-Nong Gong
 */
public class DependentProductChainSubstitutionModel extends ProductChainSubstitutionModel {

    public DependentProductChainSubstitutionModel(String name, List<SubstitutionModel> baseModels,
                                                  Parameter dependenceParameter) {
        this(name, baseModels, null, dependenceParameter);
    }

    public DependentProductChainSubstitutionModel(String name, List<SubstitutionModel> baseModels,
                                                  List<SiteRateModel> rateModels, Parameter dependenceParameter) {
        super(name, baseModels, rateModels);
        this.dependenceParameter = dependenceParameter;

    }

    public EigenDecomposition getEigenDecomposition() {
        synchronized (this) {
            if (updateMatrix) {
                computeKroneckerSumsAndProducts();
            }
        }
        return eigenDecomposition;
    }

    public void getInfinitesimalMatrix(double[] out) {
            getEigenDecomposition(); // Updates rate matrix if necessary
            System.arraycopy(rateMatrix, 0, out, 0, stateCount * stateCount);
        }

    private void computeKroneckerSumsAndProducts() {

        int currentStateSize = stateSizes[0];
        double[] currentRate = new double[currentStateSize * currentStateSize];
        baseModels.get(0).getInfinitesimalMatrix(currentRate);
        currentRate = scaleForProductChain(currentRate, 0);

        for (int i = 1; i < numBaseModel; i++) {
            SubstitutionModel nextModel = baseModels.get(i);
            int nextStateSize = stateSizes[i];
            double[] nextRate = new double[nextStateSize * nextStateSize];
            nextModel.getInfinitesimalMatrix(nextRate);
            nextRate = scaleForProductChain(nextRate, i);
            currentRate = KroneckerOperation.sum(currentRate, currentStateSize, nextRate, nextStateSize);
            currentStateSize *= nextStateSize;
        }

        double dependence = dependenceParameter.getParameterValue(0);

        /* put dependence parameter into anti-diagonal elements*/
            for (int j = 0; j < stateCount; j++) {
                currentRate[(j+1) * stateCount - (j+1)] = dependence;
            }

        /* put dependence parameter into co-varying sites */

            /* take nucleotide for example, A, T, C, and G as 0, 1, 2, and 3 */
            String[] letA = new String[(int) Math.sqrt(stateCount)];
            String[] letB = new String[(int) Math.sqrt(stateCount)];
            String[] letAB = new String[stateCount];
            int count = 0;
    
            for(int letI=0;letI<letA.length;letI++){
                letA[letI]=String.valueOf(count);
                letB[letI]=String.valueOf(count);
                count++;
            }

            /* generate two-state array {(00), (01), (02), (03), (10), ..., (33)} */
            int letABi=0;
            for(int letAi=0;letAi<letA.length;letAi++){
                for(int letBi=0;letBi<letB.length;letBi++){
                letAB[letABi]=letA[letAi]+" "+letB[letBi];
                letABi++;
                }
            }

            /* assign dependence parameter to co-varying sites */
            for(int j=0;j<letAB.length;j++){
                for(int k=0;k<letAB.length;k++){
                String[] Spl1 = letAB[j].split(" ");
                String[] Spl2 = letAB[k].split(" ");

                if(!Spl1[0].equals(Spl2[0]) && !Spl1[1].equals(Spl2[1])){
					/* Is co-varying site */
                    currentRate[j * stateCount + k] = dependence;
                }

                }
            }

        /* Rescale diagonal elements in currentRate */
            double diagonalSum=0;
            for (int k = 0; k < stateCount; k++) {
                for(int l=k*stateCount;l<((k+1)*stateCount);l++){
                    if(l!=(k * stateCount + k)){
                    diagonalSum += currentRate[l];
                    }
                }
                /* assign value to diagonal element */
                currentRate[k * stateCount + k] = -(diagonalSum);
                diagonalSum=0;
            }

        rateMatrix = currentRate;

        // Call standard eigendecomposition on currentRate
        decompose();

    }

    private void decompose() {

        double[][] q = new double[stateCount][stateCount];
        for (int i = 0; i < stateCount; i++) {
            System.arraycopy(rateMatrix, i * stateCount, q[i], 0, stateCount);
        }

        eigenDecomposition = eigenSystemCES.decomposeMatrix(q);
        //System.err.println("eigenSystem null? " + (eigenSystem == null ? "yes" : "no"));
        //eigenDecomposition = eigenSystem.decomposeMatrix(q);

        updateMatrix = false;
    }
    
    private final Parameter dependenceParameter;

    private final EigenSystem eigenSystemCES = new ColtEigenSystem();

}
