/*
 * CodonPartitionedRobustCounting.java
 *
 * Copyright (c) 2002-2012 Alexei Drummond, Andrew Rambaut and Marc Suchard
 *
 * This file is part of BEAST.
 * See the NOTICE file distributed with this work for additional
 * information regarding copyright ownership and licensing.
 *
 * BEAST is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as
 * published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 *
 *  BEAST is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with BEAST; if not, write to the
 * Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
 * Boston, MA  02110-1301  USA
 */

package dr.app.beagle.evomodel.substmodel;

import dr.app.beagle.evomodel.sitemodel.SiteRateModel;
import dr.app.beagle.evomodel.treelikelihood.AncestralStateBeagleTreeLikelihood;
import dr.app.beagle.evomodel.utilities.TreeTraitLogger;
import dr.evolution.datatype.Codons;
import dr.evolution.tree.NodeRef;
import dr.evolution.tree.Tree;
import dr.evolution.tree.TreeTrait;
import dr.evolution.tree.TreeTraitProvider;
import dr.evomodel.branchratemodel.BranchRateModel;
import dr.evomodel.tree.TreeModel;
import dr.inference.loggers.LogColumn;
import dr.inference.loggers.Loggable;
import dr.inference.markovjumps.StateHistory;
import dr.inference.model.AbstractModel;
import dr.inference.model.Model;
import dr.inference.model.Parameter;
import dr.inference.model.Variable;
import dr.math.MathUtils;

import java.util.ArrayList;
import java.util.List;

/**
 * @author Marc A. Suchard
 * @author Vladimir Minin
 *         <p/>
 *         A class for implementing robust counting for synonymous and nonsynonymous changes 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 CodonPartitionedRobustCounting extends AbstractModel implements TreeTraitProvider, Loggable {

    private static final boolean DEBUG = false;

    public static final String UNCONDITIONED_PREFIX = "u_";
    public static final String SITE_SPECIFIC_PREFIX = "c_";
    public static final String TOTAL_PREFIX = "total_";
    public static final String BASE_TRAIT_PREFIX = "base_";
    public static final String COMPLETE_HISTORY_PREFIX = "all_";
    public static final String UNCONDITIONED_PER_BRANCH_PREFIX = "b_u_";

//    public CodonPartitionedRobustCounting(String name, TreeModel tree,
//                                          AncestralStateBeagleTreeLikelihood[] partition,
//                                          Codons codons,
//                                          CodonLabeling codonLabeling,
//                                          boolean useUniformization) {
//        this(name, tree, partition, codons, codonLabeling, useUniformization,
//                StratifiedTraitOutputFormat.SUM_OVER_SITES, StratifiedTraitOutputFormat.SUM_OVER_SITES);
//
//    }

    public CodonPartitionedRobustCounting(String name, TreeModel tree,
                                          AncestralStateBeagleTreeLikelihood[] partition,
                                          Codons codons,
                                          CodonLabeling codonLabeling,
                                          boolean useUniformization,
                                          boolean includeExternalBranches,
                                          boolean includeInternalBranches,
                                          boolean doUnconditionalPerBranch,
                                          boolean saveCompleteHistory,
                                          StratifiedTraitOutputFormat branchFormat,
                                          StratifiedTraitOutputFormat logFormat) {
        super(name);
        this.tree = tree;
        addModel(tree);

        if (partition.length != 3) {
            throw new RuntimeException("CodonPartition models require 3 partitions");
        }

        this.partition = partition;
        this.codonLabeling = codonLabeling;
        branchRateModel = partition[0].getBranchRateModel();
        addModel(branchRateModel);

        List<SubstitutionModel> substModelsList = new ArrayList<SubstitutionModel>(3);
        List<SiteRateModel> siteRateModelsList = new ArrayList<SiteRateModel>(3);

        numCodons = partition[0].getPatternWeights().length;

        for (int i = 0; i < 3; i++) {
            substModelsList.add(partition[i].getSubstitutionModel());
            siteRateModelsList.add(partition[i].getSiteRateModel());
            if (partition[i].getPatternWeights().length != numCodons) {
                throw new RuntimeException("All sequence lengths must be equal in CodonPartitionedRobustCounting");
            }
        }

        this.saveCompleteHistory = saveCompleteHistory;

        productChainModel =
                new ProductChainSubstitutionModel("codonLabeling", substModelsList, siteRateModelsList);
        addModel(productChainModel);

        this.useUniformization = useUniformization;
        if (useUniformization) {
            markovJumps = new UniformizedSubstitutionModel(productChainModel);
            ((UniformizedSubstitutionModel) markovJumps).setSaveCompleteHistory(saveCompleteHistory);
        } else {
            markovJumps = new MarkovJumpsSubstitutionModel(productChainModel);
        }
        double[] synRegMatrix = CodonLabeling.getRegisterMatrix(codonLabeling, codons, true);
        markovJumps.setRegistration(synRegMatrix);

        condMeanMatrix = new double[64 * 64];

        this.branchFormat = branchFormat;
        this.logFormat = logFormat;

        computedCounts = new double[tree.getNodeCount()][]; // TODO Temporary until there exists a helper class

        this.includeExternalBranches = includeExternalBranches;
        this.includeInternalBranches = includeInternalBranches;
        this.doUnconditionedPerBranch = doUnconditionalPerBranch;

        setupTraits();
    }

    public double[] getUnconditionalCountsForBranch(NodeRef child) {
        if (!unconditionsPerBranchKnown) {
            computeAllUnconditionalCountsPerBranch();
            unconditionsPerBranchKnown = true;
        }
        return unconditionedCountsPerBranch[child.getNumber()];
    }

    public double[] getExpectedCountsForBranch(NodeRef child) { // TODO This function will implement TraitProvider
        if (!countsKnown) {
            computeAllExpectedCounts();
        }
        return computedCounts[child.getNumber()];
    }

    private void computeAllExpectedCounts() {
        for (int i = 0; i < tree.getNodeCount(); i++) {
            NodeRef child = tree.getNode(i);
            if (!tree.isRoot(child)) {
                computedCounts[child.getNumber()] = computeExpectedCountsForBranch(child);
            }
        }
        countsKnown = true;
    }

    private double[] computeExpectedCountsForBranch(NodeRef child) {

        // Get child node reconstructed sequence
        final int[] childSeq0 = partition[0].getStatesForNode(tree, child);
        final int[] childSeq1 = partition[1].getStatesForNode(tree, child);
        final int[] childSeq2 = partition[2].getStatesForNode(tree, child);

        // Get parent node reconstructed sequence
        final NodeRef parent = tree.getParent(child);
        final int[] parentSeq0 = partition[0].getStatesForNode(tree, parent);
        final int[] parentSeq1 = partition[1].getStatesForNode(tree, parent);
        final int[] parentSeq2 = partition[2].getStatesForNode(tree, parent);

        double branchRateTime = branchRateModel.getBranchRate(tree, child) * tree.getBranchLength(child);

        double[] count = new double[numCodons];

        if (!useUniformization) {
            markovJumps.computeCondStatMarkovJumps(branchRateTime, condMeanMatrix);
        } else {
            // Fill condMeanMatrix with transition probabilities
            markovJumps.getSubstitutionModel().getTransitionProbabilities(branchRateTime, condMeanMatrix);
        }

        for (int i = 0; i < numCodons; i++) {

            // Construct this child and parent codon

            final int childState = getCanonicalState(childSeq0[i], childSeq1[i], childSeq2[i]);
            final int parentState = getCanonicalState(parentSeq0[i], parentSeq1[i], parentSeq2[i]);

//            final int vChildState = getVladimirState(childSeq0[i], childSeq1[i], childSeq2[i]);
//            final int vParentState = getVladimirState(parentSeq0[i], parentSeq1[i], parentSeq2[i]);

            final double codonCount;

            if (!useUniformization) {
                codonCount = condMeanMatrix[parentState * 64 + childState];
            } else {
                codonCount = ((UniformizedSubstitutionModel) markovJumps).computeCondStatMarkovJumps(
                        parentState,
                        childState,
                        branchRateTime,
                        condMeanMatrix[parentState * 64 + childState]
                );
            }


            if (useUniformization && saveCompleteHistory) {
                UniformizedSubstitutionModel usModel = (UniformizedSubstitutionModel) markovJumps;

                StateHistory history = usModel.getStateHistory();

                // Only report syn or nonsyn changes
                double[] register = usModel.getRegistration();
                history = history.filterChanges(register);

                int historyCount = history.getNumberOfJumps();
                if (historyCount > 0) {
                    double parentTime = tree.getNodeHeight(tree.getParent(child));
                    double childTime = tree.getNodeHeight(child);
                    history.rescaleTimesOfEvents(parentTime, childTime);

                    int n = history.getNumberOfJumps();
                    // MAS may have broken the next line
                    String hstring = "{" + (i + 1) + "," + history.toStringChanges(usModel.dataType) + "}";
            if (DEBUG) {
                        System.err.println("site " + (i + 1) + " : "
                                + history.getNumberOfJumps()
                                + " : "
                                + history.toStringChanges(usModel.dataType)
                                + " " + codonLabeling.getText());
                    }

                    if (completeHistoryPerNode == null) {
                        completeHistoryPerNode = new String[tree.getNodeCount()][numCodons];
            }
//                    completeHistoryPerNode[child.getNumber()][i] = usModel.getCompleteHistory(parentTime, childTime); // TODO Should reformat
                    completeHistoryPerNode[child.getNumber()][i] = hstring;
                }
            }

            count[i] = codonCount;
        }

        return count;
    }

    private void setupTraits() {

        TreeTrait baseTrait = new TreeTrait.DA() {

            public String getTraitName() {
                return BASE_TRAIT_PREFIX + codonLabeling.getText();
            }

            public Intent getIntent() {
                return Intent.BRANCH;
            }

            public double[] getTrait(Tree tree, NodeRef node) {
                return getExpectedCountsForBranch(node);
            }

            public boolean getLoggable() {
                return false;
            }
        };

        if (saveCompleteHistory) {
            TreeTrait stringTrait = new TreeTrait.SA() {

                public String getTraitName() {
                    return COMPLETE_HISTORY_PREFIX + codonLabeling.getText();
                }

                public Intent getIntent() {
                    return Intent.BRANCH;
                }

                public String[] getTrait(Tree tree, NodeRef node) {
                    getExpectedCountsForBranch(node); // Lazy simulation of complete histories
                    List<String> events = new ArrayList<String>();
                    for (int i = 0; i < numCodons; i++) {
                        if (completeHistoryPerNode[node.getNumber()][i] != null) {
                            events.add(completeHistoryPerNode[node.getNumber()][i]);
                        }
                    }
                    String[] array = new String[events.size()];
                    events.toArray(array);
                    return array;
                }

                public boolean getLoggable() {
                    return true;
                }
            };
            treeTraits.addTrait(stringTrait);
        }

        TreeTrait unconditionedSum;
        if (!TRIAL) {
            unconditionedSum = new TreeTrait.D() {
                public String getTraitName() {
                    return UNCONDITIONED_PREFIX + codonLabeling.getText();
                }

                public Intent getIntent() {
                    return Intent.WHOLE_TREE;
                }

                public Double getTrait(Tree tree, NodeRef node) {
                    return getUnconditionedTraitValue();
                }

                public boolean getLoggable() {
                    return false;
                }
            };
        } else {
            unconditionedSum = new TreeTrait.DA() {
                public String getTraitName() {
                    return UNCONDITIONED_PREFIX + codonLabeling.getText();
                }

                public Intent getIntent() {
                    return Intent.WHOLE_TREE;
                }

                public double[] getTrait(Tree tree, NodeRef node) {
                    return getUnconditionedTraitValues();
                }

                public boolean getLoggable() {
                    return false;
                }
            };
        }

        TreeTrait sumOverTreeTrait = new TreeTrait.SumOverTreeDA(
                SITE_SPECIFIC_PREFIX + codonLabeling.getText(),
                baseTrait,
                includeExternalBranches,
                includeInternalBranches) {
            @Override
            public boolean getLoggable() {
                return false;
            }
        };

        // This should be the default output in tree logs
        TreeTrait sumOverSitesTrait = new TreeTrait.SumAcrossArrayD(
                codonLabeling.getText(),
                baseTrait) {
            @Override
            public boolean getLoggable() {
                return true;
            }
        };

        // This should be the default output in columns logs
        TreeTrait sumOverSitesAndTreeTrait = new TreeTrait.SumOverTreeD(
                TOTAL_PREFIX + codonLabeling.getText(),
                sumOverSitesTrait,
                includeExternalBranches,
                includeInternalBranches) {
            @Override
            public boolean getLoggable() {
                return true;
            }
        };


        treeTraitLogger = new TreeTraitLogger(
                tree,
                new TreeTrait[]{sumOverSitesAndTreeTrait}
        );

        treeTraits.addTrait(baseTrait);
        treeTraits.addTrait(unconditionedSum);
        treeTraits.addTrait(sumOverSitesTrait);
        treeTraits.addTrait(sumOverTreeTrait);
        treeTraits.addTrait(sumOverSitesAndTreeTrait);

        if (doUnconditionedPerBranch) {
            TreeTrait unconditionedBase = new TreeTrait.DA() {

                public String getTraitName() {
                    return UNCONDITIONED_PER_BRANCH_PREFIX + codonLabeling.getText();
                }

                public Intent getIntent() {
                    return Intent.BRANCH;
                }

                public double[] getTrait(Tree tree, NodeRef node) {
                    return getUnconditionalCountsForBranch(node);
                }

                public boolean getLoggable() {
                    return false;   // TODO Should be switched to true to log unconditioned values per branch
                }
            };

            TreeTrait sumUnconditionedOverSitesTrait = new TreeTrait.SumAcrossArrayD(
                    UNCONDITIONED_PER_BRANCH_PREFIX + codonLabeling.getText(),
                    unconditionedBase) {
                @Override
                public boolean getLoggable() {
                    return true;
                }
            };

            treeTraits.addTrait(unconditionedBase);
            treeTraits.addTrait(sumUnconditionedOverSitesTrait);
        }
    }

    public TreeTrait[] getTreeTraits() {
        return treeTraits.getTreeTraits();
    }

    public TreeTrait getTreeTrait(String key) {
        return treeTraits.getTreeTrait(key);
    }

    private int getCanonicalState(int i, int j, int k) {
        return i * 16 + j * 4 + k;
    }

//    private int getVladimirState(int i, int j, int k) {
//        if (i == 1) i = 2;
//        else if (i == 2) i = 1;
//
//        if (j == 1) j = 2;
//        else if (j == 2) j = 1;
//
//        if (k == 1) k = 2;
//        else if (k == 2) k = 1;
//
//        return i * 16 + j * 4 + k + 1;
//    }

    public LogColumn[] getColumns() {
        return treeTraitLogger.getColumns();
    }

    public int getDimension() {
        return numCodons;
    }

    private void computeAllUnconditionalCountsPerBranch() {
        if (unconditionedCountsPerBranch == null) {
            unconditionedCountsPerBranch = new double[tree.getNodeCount()][numCodons];
        }
        double[] rootDistribution = productChainModel.getFrequencyModel().getFrequencies();
        for (int i = 0; i < tree.getNodeCount(); i++) {
            NodeRef node = tree.getNode(i);
            if (!tree.isRoot(node)) {
                final double expectedLength = getExpectedBranchLength(node);
                fillInUnconditionalTraitValues(expectedLength, rootDistribution, unconditionedCountsPerBranch[node.getNumber()]);
            }
        }
    }

    private void computeUnconditionedTraitValues() {
        if (unconditionedCounts == null) {
            unconditionedCounts = new double[numCodons];
        }
        final double treeLength = getExpectedTreeLength();
        double[] rootDistribution = productChainModel.getFrequencyModel().getFrequencies();
//        final int stateCount = 64;
//        double[] lambda = new double[stateCount * stateCount];
//        productChainModel.getInfinitesimalMatrix(lambda);
//        for (int i = 0; i < numCodons; i++) {
//            final int startingState = MathUtils.randomChoicePDF(rootDistribution);
//            StateHistory history = StateHistory.simulateUnconditionalOnEndingState(
//                    0.0,
//                    startingState,
//                    treeLength,
//                    lambda,
//                    stateCount
//            );
//            unconditionedCounts[i] = markovJumps.getProcessForSimulant(history);
//        }
        fillInUnconditionalTraitValues(treeLength, rootDistribution, unconditionedCounts);
    }

    private void fillInUnconditionalTraitValues(double expectedLength, double[] freq, double[] out) {
        final int stateCount = 64;
        double[] lambda = new double[stateCount * stateCount];
        productChainModel.getInfinitesimalMatrix(lambda);
        for (int i = 0; i < numCodons; i++) {
            final int startingState = MathUtils.randomChoicePDF(freq);
            StateHistory history = StateHistory.simulateUnconditionalOnEndingState(
                    0.0,
                    startingState,
                    expectedLength,
                    lambda,
                    stateCount
            );
            out[i] = markovJumps.getProcessForSimulant(history);
        }
    }

    private double[] getUnconditionedTraitValues() {
        if (!unconditionsKnown) {
            computeUnconditionedTraitValues();
            unconditionsKnown = true;
        }
        return unconditionedCounts;
    }

    public Double getUnconditionedTraitValue() {
        if (!TRIAL) {
            return markovJumps.getMarginalRate() * getExpectedTreeLength();
        } else {
            final double treeLength = getExpectedTreeLength();
            double[] rootDistribution = productChainModel.getFrequencyModel().getFrequencies();
            final int startingState = MathUtils.randomChoicePDF(rootDistribution);
            final int stateCount = 64;
            double[] lambda = new double[stateCount * stateCount];
            productChainModel.getInfinitesimalMatrix(lambda);
            StateHistory history = StateHistory.simulateUnconditionalOnEndingState(
                    0.0,
                    startingState,
                    treeLength,
                    lambda,
                    stateCount
            );
            return markovJumps.getProcessForSimulant(history);
        }
    }

    private double getExpectedBranchLength(NodeRef node) {
        return branchRateModel.getBranchRate(tree, node) * tree.getBranchLength(node);
    }

    private double getExpectedTreeLength() {
        double expectedTreeLength = 0;
        if (includeExternalBranches) {
            for (int i = 0; i < tree.getExternalNodeCount(); i++) {
                NodeRef node = tree.getExternalNode(i);
                expectedTreeLength += getExpectedBranchLength(node);
            }
        }
        if (includeInternalBranches) {
            for (int i = 0; i < tree.getInternalNodeCount(); i++) {
                NodeRef node = tree.getInternalNode(i);
                if (!tree.isRoot(node)) {
                    expectedTreeLength += getExpectedBranchLength(node);
                }
            }
        }
        return expectedTreeLength;
    }

    protected void handleModelChangedEvent(Model model, Object object, int index) {
        countsKnown = false;
        unconditionsKnown = false;
        unconditionsPerBranchKnown = false;
    }

    protected void handleVariableChangedEvent(Variable variable, int index, Parameter.ChangeType type) {
        countsKnown = false;
        unconditionsKnown = false;
    }

    protected void storeState() {
        // Do nothing
    }

    protected void restoreState() {
        countsKnown = false;
        unconditionsKnown = false;
        unconditionsPerBranchKnown = false;
    }

    protected void acceptState() {
        // Do nothing
    }

    private final AncestralStateBeagleTreeLikelihood[] partition;
    private final MarkovJumpsSubstitutionModel markovJumps;
    private final boolean useUniformization;
    private final BranchRateModel branchRateModel;
    private final ProductChainSubstitutionModel productChainModel;

    private final CodonLabeling codonLabeling;
    private final Tree tree;

    private final StratifiedTraitOutputFormat branchFormat;
    private final StratifiedTraitOutputFormat logFormat;

    private final double[] condMeanMatrix;

    private int numCodons;

    private boolean countsKnown = false;
    private boolean unconditionsKnown = false;
    private boolean unconditionsPerBranchKnown = false;
    private double[] unconditionedCounts;
    private double[][] unconditionedCountsPerBranch;
    private double[][] computedCounts; // TODO Temporary storage until generic TreeTraitProvider/Helpers are finished

    private String[][] completeHistoryPerNode;

    protected Helper treeTraits = new Helper();
    protected TreeTraitLogger treeTraitLogger;

    private final boolean includeExternalBranches;
    private final boolean includeInternalBranches;
    private final boolean doUnconditionedPerBranch;

    private static final boolean TRIAL = true;

    private boolean saveCompleteHistory = false;

}
