/*
 * GeneralLikelihoodCore.java
 *
 * Copyright (C) 2002-2006 Alexei Drummond and Andrew Rambaut
 *
 * 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.evomodel.treelikelihood;

import java.math.BigDecimal;
import java.math.MathContext;

/**
 * GeneralLikelihoodCore - An implementation of LikelihoodCore for any data
 *
 * @version $Id: GeneralLikelihoodCore.java,v 1.28 2006/08/31 14:57:24 rambaut Exp $
 *
 * @author Andrew Rambaut
 */

public class ArbitraryPrecisionLikelihoodCore implements LikelihoodCore {

    private MathContext precision;

    private int stateCount;
    private int nodeCount;
    private int patternCount;
    private int partialsSize;
    private int matrixSize;
    private int matrixCount;

    private boolean integrateCategories;

    private BigDecimal[][][] partials;

    private int[][] states;

    private BigDecimal[][][] matrices;

    private int[] currentMatricesIndices;
    private int[] storedMatricesIndices;
    private int[] currentPartialsIndices;
    private int[] storedPartialsIndices;


	/**
	 * Constructor
	 *
	 * @param stateCount number of states
	 */
	public ArbitraryPrecisionLikelihoodCore(int stateCount, int precision) {
        this.stateCount = stateCount;
        this.precision = new MathContext(precision);
    }

    /**
     * initializes partial likelihood arrays.
     *
     * @param nodeCount           the number of nodes in the tree
     * @param patternCount        the number of patterns
     * @param matrixCount         the number of matrices (i.e., number of categories)
     * @param integrateCategories whether sites are being integrated over all matrices
     */
    public void initialize(int nodeCount, int patternCount, int matrixCount, boolean integrateCategories) {

        this.nodeCount = nodeCount;
        this.patternCount = patternCount;
        this.matrixCount = matrixCount;

        this.integrateCategories = integrateCategories;

        if (integrateCategories) {
            partialsSize = patternCount * stateCount * matrixCount;
        } else {
            partialsSize = patternCount * stateCount;
        }

        partials = new BigDecimal[2][nodeCount][];

        currentMatricesIndices = new int[nodeCount];
        storedMatricesIndices = new int[nodeCount];

        currentPartialsIndices = new int[nodeCount];
        storedPartialsIndices = new int[nodeCount];

        states = new int[nodeCount][];

        for (int i = 0; i < nodeCount; i++) {
            partials[0][i] = null;
            partials[1][i] = null;

            states[i] = null;
        }

        matrixSize = stateCount * stateCount;

        matrices = new BigDecimal[2][nodeCount][matrixCount * matrixSize];
    }

    /**
     * cleans up and deallocates arrays.
     */
    public void finalize() throws java.lang.Throwable  {
        super.finalize();

        nodeCount = 0;
        patternCount = 0;
        matrixCount = 0;

        partials = null;
        currentPartialsIndices = null;
        storedPartialsIndices = null;
        states = null;
        matrices = null;
        currentMatricesIndices = null;
        storedMatricesIndices = null;
    }

    /**
     * Allocates partials for a node
     */
    public void createNodePartials(int nodeIndex) {

        this.partials[0][nodeIndex] = new BigDecimal[partialsSize];
        this.partials[1][nodeIndex] = new BigDecimal[partialsSize];
    }

    /**
     * Sets partials for a node
     */
    public void setNodePartials(int nodeIndex, double[] partials) {

        if (this.partials[0][nodeIndex] == null) {
            createNodePartials(nodeIndex);
        }
        if (partials.length < partialsSize) {
            int k = 0;
            for (int i = 0; i < matrixCount; i++) {
                for (int j = 0; j < partials.length; j++) {
                    this.partials[0][nodeIndex][k] = new BigDecimal(partials[j], precision);
                    k++;
                }
            }
        } else {
            for (int j = 0; j < partials.length; j++) {
                this.partials[0][nodeIndex][j] = new BigDecimal(partials[j], precision);
            }
        }
    }

    /**
     * Allocates states for a node
     */
    public void createNodeStates(int nodeIndex) {

        this.states[nodeIndex] = new int[patternCount];
    }

    /**
     * Sets states for a node
     */
    public void setNodeStates(int nodeIndex, int[] states) {

        if (this.states[nodeIndex] == null) {
            createNodeStates(nodeIndex);
        }
        System.arraycopy(states, 0, this.states[nodeIndex], 0, patternCount);
    }

    /**
     * Gets states for a node
     */
    public void getNodeStates(int nodeIndex, int[] states) {
        System.arraycopy(this.states[nodeIndex], 0, states, 0, patternCount);
    }

    public void setNodeMatrixForUpdate(int nodeIndex) {
        currentMatricesIndices[nodeIndex] = 1 - currentMatricesIndices[nodeIndex];

    }

    /**
     * Sets probability matrix for a node
     */
    public void setNodeMatrix(int nodeIndex, int matrixIndex, double[] matrix) {

        for (int j = 0; j < matrixSize; j++) {
            matrices[currentMatricesIndices[nodeIndex]][nodeIndex][(matrixIndex * matrixSize) + j] = new BigDecimal(matrix[j], precision);
        }
    }

    /**
     * Gets probability matrix for a node
     */
    public void getNodeMatrix(int nodeIndex, int matrixIndex, double[] matrix) {
        for (int j = 0; j < matrixSize; j++) {
            matrix[j] = matrices[currentMatricesIndices[nodeIndex]][nodeIndex][(matrixIndex * matrixSize) + j].doubleValue();
        }
    }

    public void setNodePartialsForUpdate(int nodeIndex) {
        currentPartialsIndices[nodeIndex] = 1 - currentPartialsIndices[nodeIndex];
    }

    /**
     * Get the currently updating node partials for direct access
     *
     * @param nodeIndex
     */
    public void setCurrentNodePartials(int nodeIndex, double[] partials) {
        throw new UnsupportedOperationException("setCurrentNodePartials is not supported by ArbitraryPrecisionLikelihoodCore");
    }

    /**
     * Calculates partial likelihoods at a node.
     *
     * @param nodeIndex1 the 'child 1' node
     * @param nodeIndex2 the 'child 2' node
     * @param nodeIndex3 the 'parent' node
     */
    public void calculatePartials(int nodeIndex1, int nodeIndex2, int nodeIndex3) {
        if (states[nodeIndex1] != null) {
            if (states[nodeIndex2] != null) {
                calculateStatesStatesPruning(
                        states[nodeIndex1], matrices[currentMatricesIndices[nodeIndex1]][nodeIndex1],
                        states[nodeIndex2], matrices[currentMatricesIndices[nodeIndex2]][nodeIndex2],
                        partials[currentPartialsIndices[nodeIndex3]][nodeIndex3]);
            } else {
                calculateStatesPartialsPruning(states[nodeIndex1], matrices[currentMatricesIndices[nodeIndex1]][nodeIndex1],
                        partials[currentPartialsIndices[nodeIndex2]][nodeIndex2], matrices[currentMatricesIndices[nodeIndex2]][nodeIndex2],
                        partials[currentPartialsIndices[nodeIndex3]][nodeIndex3]);
            }
        } else {
            if (states[nodeIndex2] != null) {
                calculateStatesPartialsPruning(states[nodeIndex2], matrices[currentMatricesIndices[nodeIndex2]][nodeIndex2],
                        partials[currentPartialsIndices[nodeIndex1]][nodeIndex1], matrices[currentMatricesIndices[nodeIndex1]][nodeIndex1],
                        partials[currentPartialsIndices[nodeIndex3]][nodeIndex3]);
            } else {
                calculatePartialsPartialsPruning(partials[currentPartialsIndices[nodeIndex1]][nodeIndex1], matrices[currentMatricesIndices[nodeIndex1]][nodeIndex1],
                        partials[currentPartialsIndices[nodeIndex2]][nodeIndex2], matrices[currentMatricesIndices[nodeIndex2]][nodeIndex2],
                        partials[currentPartialsIndices[nodeIndex3]][nodeIndex3]);
            }
        }
    }

	/**
	 * Calculates partial likelihoods at a node when both children have states.
	 */
	private void calculateStatesStatesPruning(int[] states1, BigDecimal[] matrices1,
												int[] states2, BigDecimal[] matrices2,
												BigDecimal[] partials3)
	{
		int v = 0;

		for (int l = 0; l < matrixCount; l++) {

			for (int k = 0; k < patternCount; k++) {

				int state1 = states1[k];
				int state2 = states2[k];

				int w = l * matrixSize;

                if (state1 < stateCount && state2 < stateCount) {

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

						partials3[v] = matrices1[w + state1].multiply(matrices2[w + state2], precision);

						v++;
						w += stateCount;
					}

				} else if (state1 < stateCount) {
					// child 2 has a gap or unknown state so treat it as unknown

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

						partials3[v] = matrices1[w + state1];

						v++;
						w += stateCount;
					}
				} else if (state2 < stateCount) {
					// child 2 has a gap or unknown state so treat it as unknown

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

						partials3[v] = matrices2[w + state2];

						v++;
						w += stateCount;
					}
				} else {
					// both children have a gap or unknown state so set partials to 1

					for (int j = 0; j < stateCount; j++) {
						partials3[v] = BigDecimal.ONE;
						v++;
					}
				}
			}
		}
	}

	/**
	 * Calculates partial likelihoods at a node when one child has states and one has partials.
	 */
	private void calculateStatesPartialsPruning(	int[] states1, BigDecimal[] matrices1,
													BigDecimal[] partials2, BigDecimal[] matrices2,
													BigDecimal[] partials3)
	{

		BigDecimal sum, tmp;

		int u = 0;
		int v = 0;

		for (int l = 0; l < matrixCount; l++) {
			for (int k = 0; k < patternCount; k++) {

				int state1 = states1[k];

                int w = l * matrixSize;

				if (state1 < stateCount) {


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

						tmp = matrices1[w + state1];

						sum = BigDecimal.ZERO;
						for (int j = 0; j < stateCount; j++) {
							sum = sum.add(matrices2[w].multiply(partials2[v + j], precision), precision);
							w++;
						}

						partials3[u] = tmp.multiply(sum, precision);
						u++;
					}

					v += stateCount;
				} else {
					// Child 1 has a gap or unknown state so don't use it

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

                        sum = BigDecimal.ZERO;
						for (int j = 0; j < stateCount; j++) {
							sum = sum.add(matrices2[w].multiply(partials2[v + j], precision), precision);
							w++;
						}

						partials3[u] = sum;
						u++;
					}

					v += stateCount;
				}
			}
		}
	}

	/**
	 * Calculates partial likelihoods at a node when both children have partials.
	 */
	private void calculatePartialsPartialsPruning(BigDecimal[] partials1, BigDecimal[] matrices1,
													BigDecimal[] partials2, BigDecimal[] matrices2,
													BigDecimal[] partials3)
	{
		BigDecimal sum1, sum2;

		int u = 0;
		int v = 0;

		for (int l = 0; l < matrixCount; l++) {

			for (int k = 0; k < patternCount; k++) {

                int w = l * matrixSize;

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

					sum1 = sum2 = BigDecimal.ZERO;

					for (int j = 0; j < stateCount; j++) {
						sum1 = sum1.add(matrices1[w].multiply(partials1[v + j], precision), precision);
						sum2 = sum2.add(matrices2[w].multiply(partials2[v + j], precision), precision);

						w++;
					}

					partials3[u] = sum1.multiply(sum2, precision);
					u++;
				}
				v += stateCount;
			}
		}
	}

    /**
     * Calculates partial likelihoods at a node.
     *
     * @param nodeIndex1 the 'child 1' node
     * @param nodeIndex2 the 'child 2' node
     * @param nodeIndex3 the 'parent' node
     * @param matrixMap  a map of which matrix to use for each pattern (can be null if integrating over categories)
     */
    public void calculatePartials(int nodeIndex1, int nodeIndex2, int nodeIndex3, int[] matrixMap) {
         throw new UnsupportedOperationException("calculatePartials(int nodeIndex1, int nodeIndex2, int nodeIndex3, int[] matrixMap) is not implemented in this likelihood core");
    }


    public void integratePartials(int nodeIndex, double[] proportions, double[] outPartials) {
        throw new UnsupportedOperationException("integratePartials(int nodeIndex, double[] proportions, double[] outPartials) is not implemented in this likelihood core");
    }

    public void integratePartials(int nodeIndex, double[] proportions, BigDecimal[] outPartials) {
        BigDecimal[] prop = new BigDecimal[proportions.length];
        for (int i = 0; i < proportions.length; i++) {
            prop[i] = new BigDecimal(proportions[i], precision);
        }
        calculateIntegratePartials(partials[currentPartialsIndices[nodeIndex]][nodeIndex], prop, outPartials);
    }

	/**
	 * Integrates partials across categories.
     * @param inPartials the array of partials to be integrated
	 * @param proportions the proportions of sites in each category
	 * @param outPartials an array into which the partials will go
	 */
	private void calculateIntegratePartials(BigDecimal[] inPartials, BigDecimal[] proportions, BigDecimal[] outPartials)
	{

		int u = 0;
		int v = 0;
		for (int k = 0; k < patternCount; k++) {

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

				outPartials[u] = inPartials[v].multiply(proportions[0], precision);
				u++;
				v++;
			}
		}


		for (int l = 1; l < matrixCount; l++) {
			u = 0;

			for (int k = 0; k < patternCount; k++) {

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

					outPartials[u] = outPartials[u].add(inPartials[v].multiply(proportions[l], precision), precision);
					u++;
					v++;
				}
			}
		}
	}

    /**
     * Calculates patten log likelihoods at a node.
     * @param partials the partials used to calculate the likelihoods
     * @param frequencies an array of state frequencies
     * @param outLogLikelihoods an array into which the likelihoods will go
     */
    public void calculateLogLikelihoods(BigDecimal[] partials, double[] frequencies, double[] outLogLikelihoods)
    {
        BigDecimal[] freqs = new BigDecimal[frequencies.length];
        for (int i = 0; i < freqs.length; i++) {
            freqs[i] = new BigDecimal(frequencies[i], precision);
        }

        int v = 0;
        for (int k = 0; k < patternCount; k++) {

            BigDecimal sum = BigDecimal.ZERO;
            for (int i = 0; i < stateCount; i++) {

                sum = sum.add(freqs[i].multiply(partials[v], precision), precision);
                v++;
            }

            double scale = sum.scale();
            double value = sum.unscaledValue().doubleValue();

            outLogLikelihoods[k] = Math.log(value) - (scale * Math.log(10));
        }
    }

    public void calculateLogLikelihoods(double[] partials, double[] frequencies, double[] outLogLikelihoods) {
        throw new UnsupportedOperationException("calculateLogLikelihoods(double[] partials, double[] frequencies, double[] outLogLikelihoods) is not implemented in this likelihood core");
    }

    public void setUseScaling(boolean useScaling) {
    }

    public double getLogScalingFactor(int pattern) {
        return 0;
    }

    /**
     * Gets the partials for a particular node.
     *
     * @param nodeIndex   the node
     * @param outPartials an array into which the partials will go
     */
    public void getPartials(int nodeIndex, double[] outPartials) {
        for (int i = 0; i < outPartials.length; i++) {
            outPartials[i] = partials[currentPartialsIndices[nodeIndex]][nodeIndex][i].doubleValue();
        }
    }

    /**
     * Store current state
     */
    public void storeState() {

        System.arraycopy(currentMatricesIndices, 0, storedMatricesIndices, 0, nodeCount);
        System.arraycopy(currentPartialsIndices, 0, storedPartialsIndices, 0, nodeCount);
    }

    /**
     * Restore the stored state
     */
    public void restoreState() {
        // Rather than copying the stored stuff back, just swap the pointers...
        int[] tmp1 = currentMatricesIndices;
        currentMatricesIndices = storedMatricesIndices;
        storedMatricesIndices = tmp1;

        int[] tmp2 = currentPartialsIndices;
        currentPartialsIndices = storedPartialsIndices;
        storedPartialsIndices = tmp2;
    }

    public void checkScaling() {
        // do nothing
    }

}