/*
 * TreeColouring.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.evolution.colouring;

import dr.evolution.tree.NodeRef;
import dr.evolution.tree.Tree;

/**
 * @author Alexei Drummond
 * @author Andrew Rambaut
 *
 * @version $Id: TreeColouring.java,v 1.8 2006/08/23 10:46:33 rambaut Exp $
 */
public class DefaultTreeColouring implements TreeColouring {

    public DefaultTreeColouring(int colourCount, Tree tree) {
        this.colourCount = colourCount;
        this.tree = tree;
        branchColourings = new DefaultBranchColouring[tree.getNodeCount()];
        colourChangeCount = 0;
    }

    /* Creates a mutable copy of the colouring (without probability assigned) */
    public DefaultTreeColouring( DefaultTreeColouring treeColouring ) {

    	if (!treeColouring.immutable) {
    		throw new RuntimeException("Attempted to create a copy of a mutable colouring");
    	}
    	this.colourCount = treeColouring.colourCount;
    	this.tree = treeColouring.tree;                                 // NOT a copy, to allow tree editing after creation
    	colourChangeCount = treeColouring.colourChangeCount;
    	branchColourings = (DefaultBranchColouring[])treeColouring.branchColourings.clone();      // Shallow copy
    	for (int i=0; i<branchColourings.length; i++) {
    		branchColourings[i] = branchColourings[i].getCopy();        // Turn into a deep copy
    	}
    }

    public int getColourCount() {
        return colourCount;
    }

    public Tree getTree() {
        return tree;
    }

    public void setBranchColouring(NodeRef node, DefaultBranchColouring branchColouring) {

    	if (immutable) throw new RuntimeException("Attempted to setBranchColouring of an immutable colouring");
        if (tree.isRoot(node)) throw new IllegalArgumentException("Can't set a BranchColouring for the root node");

        if (branchColourings[node.getNumber()] != null) {
        	colourChangeCount -= branchColourings[node.getNumber()].getColourChanges().size();
        }

        branchColourings[node.getNumber()] = branchColouring;

        colourChangeCount += branchColouring.getColourChanges().size();
    }

    /*
     * Previous version - faster, but more complicated.
     *
    public BranchColouring getBranchColouring(NodeRef node) {

    	if (tree.isRoot(node)) throw new IllegalArgumentException("Can't get a BranchColouring for the root node");

    	if (mutableBranches) {
    		return branchColourings[node.getNumber()];
    	}

    	if (branchColourings[node.getNumber()] == null) return null;

    	return branchColourings[node.getNumber()].getCopy();
    }
    */

    public void setLogProbabilityDensity(double logP) {
        if (hasProbability) throw new RuntimeException("Attempted to set probability twice");
        makeImmutable();
        hasProbability = true;
        this.logP = logP;
    }

    public void makeImmutable() {
        immutable = true;
    }

    public boolean isImmutable() {
    	return immutable;
    }

    public boolean hasProbability() {
    	return hasProbability;
    }

    public void checkColouring() {
        checkColouringAtNode(tree.getRoot());
    }

    private void checkColouringAtNode(NodeRef node) {
        int colour = getNodeColour(node);

        if (!tree.isExternal(node)) {
            for(int i = 0; i < tree.getChildCount(node); i++) {
                NodeRef child = tree.getChild(node, i);
                if (branchColourings[child.getNumber()].getParentColour() != colour) {
                    throw new RuntimeException("Colour mismatch at node " + node.toString());
                }
                checkColouringAtNode(child);
            }
        }
    }

    /**
     * @param node
     * @return the colour of the given node.
     */
    public int getNodeColour(NodeRef node) {
        if (tree.isRoot(node)) {
            NodeRef child = tree.getChild(node, 0);
            return branchColourings[child.getNumber()].getParentColour();
        } else {
            return branchColourings[node.getNumber()].getChildColour();
        }
    }

    /**
     * This method will return the colour on the ancestral branch if the given time
     * is older than the parent of the given node.
     *
     * @param node
     * @return the colour of the branch above the given node at the given (absolute) time.
     */
   /* public final int getColour(NodeRef node, double time) {

        double nodeHeight = tree.getNodeHeight(node);

        if (time < nodeHeight) throw new IllegalArgumentException("Specified time is younger than node!");
        while ((!tree.isRoot(node)) && (tree.getNodeHeight(tree.getParent(node)) < time)) {
        	node = tree.getParent(node);
        }
        if (tree.isRoot(node)) throw new IllegalArgumentException("Can't ascertain lineage colour above root!");
        int colour = getColour(node);
        List changes = getColourChanges(node);
        int size = changes.size();
        int index = 0;
        while (index < size && ((ColourChange)changes.get(index)).getTime() < time) {
            index += 1;
        }
        if (index == 0) return colour;
        return ((ColourChange)changes.get(index-1)).getColourAbove();
    }
*/

    /**
     * @param node
     * @return a branch colouring above this node
     */
    public BranchColouring getBranchColouring(NodeRef node) {
        return branchColourings[node.getNumber()];
    }


    public int getColourChangeCount() {
        return colourChangeCount;
    }

    /**
     * @return the log probability density of this colouring.
     */
    public double getLogProbabilityDensity() {
    	if (!hasProbability) {
    		throw new RuntimeException("Tree colouring has no probability density; use colouringModel.getTreeColouringWithProbability()");
    	}
        return logP;
    }

    private final Tree tree;
    private final int colourCount;
    private final DefaultBranchColouring[] branchColourings;
    private int colourChangeCount;

    private double logP = Double.NEGATIVE_INFINITY;

    private boolean immutable = false;             // true if colouring cannot be changed
    private boolean hasProbability = false;        // true if probability has been assigned (implies immutable)

}
