package dr.app.beagle.evomodel.substmodel;

import cern.colt.matrix.DoubleMatrix2D;
import cern.colt.matrix.impl.DenseDoubleMatrix2D;
import cern.colt.matrix.linalg.Algebra;
import cern.colt.matrix.linalg.LUDecomposition;
import dr.inference.model.Parameter;
import dr.inference.model.Variable;

import java.util.List;

/**
 * @author Marc A. Suchard
 * @author Vladimir Minin
 *         <p/>
 *         A class for implementing a kronecker sum of CTMC models in BEAST using BEAGLE
 *         This work is supported by NSF grant 0856099
 *         <p/>
 *         O'Brien JD, Minin VN and Suchard MA (2009) Learning to count: robust estimates for labeled distances between
 *         molecular sequences. Molecular Biology and Evolution, 26, 801-814
 */

public class MarkovModulatedFrequencyModel extends FrequencyModel {

    public MarkovModulatedFrequencyModel(String name, List<FrequencyModel> freqModels, Parameter switchingRates) {
        super(name);
        this.freqModels = freqModels;
        int freqCount = 0;
        stateCount = freqModels.get(0).getFrequencyCount();
        numBaseModel = freqModels.size();
        for (int i = 0; i < numBaseModel; i++) {
            int size = freqModels.get(i).getFrequencyCount();
            if (stateCount != size) {
                throw new RuntimeException("MarkovModulatedFrequencyModel requires all frequencies model to have the same dimension");
            }           
            freqCount += size;
        }
        totalFreqCount = freqCount;
        this.switchingRates = switchingRates;
        addVariable(switchingRates);

        if (switchingRates.getDimension() > 2 ||  numBaseModel > 2) {
            throw new RuntimeException("MarkovModulatedFrequencyModel not yet implemented for more than 2 hidden classes");
        }

        baseStationaryDistribution = new double[numBaseModel];
        storedBaseStationaryDistribution = new double[numBaseModel];
        stationaryDistributionKnown = false;

        d = new DenseDoubleMatrix2D(numBaseModel, numBaseModel);
        d.set(0, 0, 1.0);
    }

    public void setFrequency(int i, double value) {
        throw new RuntimeException("Not implemented");
    }

    public double getFrequency(int index) {
        int whichModel = index / stateCount;
        int whichState = index % stateCount;
        double relativeFreq = freqModels.get(whichModel).getFrequency(whichState);

        // Scale by stationary distribution over hidden classes
        if (numBaseModel > 1) {
            if (!stationaryDistributionKnown) {
                computeStationaryDistribution(baseStationaryDistribution);
                stationaryDistributionKnown = true;
            }
            relativeFreq *= baseStationaryDistribution[whichModel];
        }

        return relativeFreq;
    }

    private void computeStationaryDistribution(double[] statDistr) {

//        // Uses an eigendecomposition and matrix inverse
//        DoubleMatrix2D mat = new DenseDoubleMatrix2D(numBaseModel, numBaseModel);
//        int index = 0;
//        for (int i = 0; i < numBaseModel; ++i) {
//            for (int j = i + 1; j < numBaseModel; ++j) {
//                mat.set(i, j, switchingRates.getParameterValue(index));
//                index++;
//            }
//        }
//        for (int j = 0; j < numBaseModel; ++j) {
//            for (int i = j + 1; i < numBaseModel; ++i) {
//                mat.set(i, j, switchingRates.getParameterValue(index));
//                index++;
//            }
//        }
//        for (int i = 0; i < numBaseModel; ++i) {
//            double rowTotal = 0.0;
//            for (int j = 0; j < numBaseModel; ++j) {
//                if (i != j) {
//                    rowTotal += mat.get(i,j);
//                }
//            }
//            mat.set(i,i, -rowTotal);
//        }
//
//        EigenvalueDecomposition ed = new EigenvalueDecomposition(mat);
//        DoubleMatrix2D eigenVectors = ed.getV();
//        DoubleMatrix2D b = alg.mult(eigenVectors, alg.mult(d, alg.inverse(eigenVectors)));
//
//        for (int i = 0; i < numBaseModel; ++i) {
//            statDistr[i] = b.get(0,i);
//        }
//        System.err.println(new Vector(statDistr));


        // Uses an LU decomposition to solve Q^t \pi = 0 and \sum \pi_i = 1
        DoubleMatrix2D mat2 = new DenseDoubleMatrix2D(numBaseModel + 1, numBaseModel);
        int index2 = 0;
        for (int i = 0; i < numBaseModel; ++i) {
            for (int j = i + 1; j < numBaseModel; ++j) {
                mat2.set(j, i, switchingRates.getParameterValue(index2)); // Transposed
                index2++;
            }
        }
        for (int j = 0; j < numBaseModel; ++j) {
            for (int i = j + 1; i < numBaseModel; ++i) {
                mat2.set(j, i, switchingRates.getParameterValue(index2)); // Transposed
                index2++;
            }
        }
        for (int i = 0; i < numBaseModel; ++i) {
            double rowTotal = 0.0;
            for (int j = 0; j < numBaseModel; ++j) {
                if (i != j) {
                    rowTotal += mat2.get(j,i); // Transposed
                }

            }
            mat2.set(i,i, -rowTotal);
        }

        // Add row for sum-to-one constraint
        for (int i = 0; i < numBaseModel; ++i) {
            mat2.set(numBaseModel, i, 1.0);
        }

        LUDecomposition decomp = new LUDecomposition(mat2);
        DoubleMatrix2D x = new DenseDoubleMatrix2D(numBaseModel + 1, 1);
        x.set(numBaseModel, 0, 1.0);
        DoubleMatrix2D y = decomp.solve(x);
        for (int i = 0; i < numBaseModel; ++i) {
            statDistr[i] = y.get(i,0);
        }
        //System.err.println(new Vector(statDistr));              
    }

    protected void storeState() {
        System.arraycopy(baseStationaryDistribution, 0, storedBaseStationaryDistribution, 0, numBaseModel);        
        storedStationaryDistributionKnown = stationaryDistributionKnown;
    }

    protected void restoreState() {
        double[] tmp = baseStationaryDistribution;
        baseStationaryDistribution = storedBaseStationaryDistribution;
        storedBaseStationaryDistribution = tmp;

        stationaryDistributionKnown = storedStationaryDistributionKnown;
    }

    protected void handleVariableChangedEvent(Variable variable, int index, Parameter.ChangeType type) {
        if (variable == switchingRates) {           
            stationaryDistributionKnown = false;
        }
    }

    public int getFrequencyCount() {
        return totalFreqCount;
    }

    public Parameter getFrequencyParameter() {
        throw new RuntimeException("Not implemented");
    }
     
    private List<FrequencyModel> freqModels;

    private final int numBaseModel;
    private final int totalFreqCount;
    private final int stateCount;
    private final Parameter switchingRates;

    private double[] baseStationaryDistribution;
    private double[] storedBaseStationaryDistribution;

    private boolean stationaryDistributionKnown;
    private boolean storedStationaryDistributionKnown;

    private final Algebra alg = new Algebra();
    private final DoubleMatrix2D d;
}
