package beast.evolution.branchratemodel;


import beast.core.Citation;
import beast.core.Description;
import beast.core.Input;
import beast.core.parameter.IntegerParameter;
import beast.core.parameter.RealParameter;
import beast.evolution.tree.Node;
import beast.evolution.tree.Tree;
import beast.math.distributions.ParametricDistribution;

/**
 * @author Alexei Drummond
 */

@Description("Defines an uncorrelated relaxed molecular clock.")
@Citation(value = "Drummond AJ, Ho SYW, Phillips MJ, Rambaut A (2006) Relaxed Phylogenetics and Dating with Confidence. PLoS Biol 4(5): e88", DOI = "10.1371/journal.pbio.0040088")
public class UCRelaxedClockModel extends BranchRateModel.Base {

    public Input<ParametricDistribution> rateDistInput = new Input<ParametricDistribution>("distr", "the distribution governing the rates among branches. Must have mean of 1. The clock.rate parameter can be used to change the mean rate.", Input.Validate.REQUIRED);
    public Input<IntegerParameter> categoryInput = new Input<IntegerParameter>("rateCategories", "the rate categories associated with nodes in the tree for sampling of individual rates among branches.", Input.Validate.REQUIRED);
    public Input<Tree> treeInput = new Input<Tree>("tree", "the tree this relaxed clock is associated with.", Input.Validate.REQUIRED);
    public Input<Boolean> normalizeInput = new Input<Boolean>("normalize", "Whether to normalize the average rate (default false).", false);

    RealParameter meanRate;
    
    @Override
    public void initAndValidate() throws Exception {

        tree = treeInput.get();

        categories = categoryInput.get();
        int nCategoryCount = tree.getNodeCount() - 1;
        categories.setDimension(nCategoryCount);
        Integer [] iCategories = new Integer[nCategoryCount];
        for (int i = 0; i < nCategoryCount; i++) {
        	iCategories[i] = i;
        }
        IntegerParameter other = new IntegerParameter(iCategories);
        categories.assignFromWithoutID(other);
        categories.setLower(0);
        categories.setUpper(categories.getDimension() - 1);
        
        distribution = rateDistInput.get();

        rates = new double[categories.getDimension()];
        storedRates = new double[categories.getDimension()];
        for (int i = 0; i < rates.length; i++) {
            rates[i] = distribution.inverseCumulativeProbability((i + 0.5) / rates.length);
        }
        System.arraycopy(rates, 0, storedRates, 0, rates.length);
        normalize = normalizeInput.get();

        meanRate = meanRateInput.get();
        if (meanRate == null) {
        	meanRate = new RealParameter("1.0");
        }
    }

    public double getRateForBranch(Node node) {
    	if (node.isRoot()) {
        	// root has no rate
    		return 1;
    	}
    	
        if (recompute) {
            prepare();
            recompute = false;
        }

        int nodeNumber = node.getNr();

        if (nodeNumber == categories.getDimension()) {
        	// root node has nr less than #categories, so use that nr
    		nodeNumber = node.getTree().getRoot().getNr();
        }

        int rateCategory = categories.getValue(nodeNumber);

        return rates[rateCategory] * scaleFactor * meanRate.getValue();
    }

    // compute scale factor

    private void computeFactor() {

        //scale mean rate to 1.0 or separate parameter

        double treeRate = 0.0;
        double treeTime = 0.0;

        //normalizeBranchRateTo = 1.0;
        for (int i = 0; i < tree.getNodeCount(); i++) {
            Node node = tree.getNode(i);
            if (!node.isRoot()) {
                int rateCategory = categories.getValue(node.getNr());
                treeRate += rates[rateCategory] * node.getLength();
                treeTime += node.getLength();

                //System.out.println("rates and time\t" + rates[rateCategory] + "\t" + node.getLength());
            }
        }
        //treeRate /= treeTime;

        scaleFactor = normalizeBranchRateTo / (treeRate / treeTime);


        //System.out.println("scaleFactor\t\t\t\t\t" + scaleFactor);
    }


    private void prepare() {

        //System.out.println("prepare");

        categories = (IntegerParameter) categoryInput.get();

        distribution = rateDistInput.get();

        tree = treeInput.get();

        rates = new double[categories.getDimension()];
        try {
	        for (int i = 0; i < rates.length; i++) {
	            rates[i] = distribution.inverseCumulativeProbability((i + 0.5) / rates.length);
	        }
        } catch (Exception e) {
        	// Exception due to distribution not having  inverseCumulativeProbability implemented.
        	// This should already been caught at initAndValidate()
        	e.printStackTrace();
        	System.exit(0);
        }

        if (normalize) computeFactor();
    }

    @Override
    protected boolean requiresRecalculation() {
    	recompute = false;

//        if (treeInput.get().somethingIsDirty()) {
//        	recompute = true;
//            return true;
//        }
        // rateDistInput cannot be dirty?!?
        if (rateDistInput.get().isDirtyCalculation()) {
        	recompute = true;
        	return true;
        }
        // NOT processed as trait on the tree, so DO mark as dirty
        if (categoryInput.get().somethingIsDirty()) {
        	//recompute = true;
        	return true;
        }
        if (meanRate.somethingIsDirty()) {
        	return true;
        }
        
        return recompute;
    }

    @Override
    public void store() {
        System.arraycopy(rates, 0, storedRates, 0, rates.length);
    	super.store();
    }
    @Override
    public void restore() {
    	double [] tmp = rates;
    	rates = storedRates;
    	storedRates = tmp;
    	super.restore();
    }
    
    ParametricDistribution distribution;
    IntegerParameter categories;
    Tree tree;

    private boolean normalize = false;
    private boolean recompute = true;
    private double normalizeBranchRateTo = Double.NaN;
    private double[] rates;
    private double[] storedRates;
    private double scaleFactor = 1.0;

}

@@ -3431,47 +3431,8 @@
 0;%0A%0A
-        //normalizeBranchRateTo = 1.0;%0A
     
@@ -3928,37 +3928,19 @@
 actor = 
-normalizeBranchRateTo
+1.0
  / (tree
@@ -5898,63 +5898,8 @@
 ue;%0A
-    private double normalizeBranchRateTo = Double.NaN;%0A
     

solve Issue 1: UCrelaxed normalize bug