package dr.evomodel.speciation;

import dr.evolution.tree.NodeRef;
import dr.evolution.tree.Tree;
import dr.inference.model.Parameter;
import static org.apache.commons.math.special.Gamma.logGamma;

/**
 * Birth Death model based on Gernhard 2008  "The conditioned reconstructed process"
 * doi:10.1016/j.jtbi.2008.04.005 (http://dx.doi.org/10.1016/j.jtbi.2008.04.005)
 * <p/>
 * This derivation conditions directly on fixed N taxa.
 * <p/>
 * The inference is directly on b-d (strictly positive) and d/b (constrained in [0,1))
 * <p/>
 * Vefified using simulated trees generated by Klass tree sample. (http://www.klaashartmann.com/treesample/)
 *
 * Sampling proportion not verified via simulation. Proportion set by default to 1, an assignment which makes the expressions
 * identical to the ones pre change.
 *
 *
 * @author joseph
 *         Date: 24/02/2008
 */
public class BirthDeathGernhard08Model extends SpeciationModel {

    public static final String BIRTH_DEATH_MODEL = "birthDeathModel";

    Parameter relativeDeathRateParameter; /** mu/lambda **/
    Parameter birthDiffRateParameter; /** lambda - mu **/
    Parameter sampleProbability; /** rho **/

    public BirthDeathGernhard08Model(Parameter birthDiffRateParameter,
                                     Parameter relativeDeathRateParameter,
                                     Parameter sampleProbability,
                                     Type units) {

        this(BIRTH_DEATH_MODEL, birthDiffRateParameter, relativeDeathRateParameter, sampleProbability, units);
    }

    BirthDeathGernhard08Model(String modelName,
                              Parameter birthDiffRateParameter, Parameter relativeDeathRateParameter, Parameter sampleProbability,
                              Type units) {

        super(modelName, units);

        this.birthDiffRateParameter = birthDiffRateParameter;
        addParameter(birthDiffRateParameter);
        birthDiffRateParameter.addBounds(new Parameter.DefaultBounds(Double.POSITIVE_INFINITY, 0.0, 1));

        this.relativeDeathRateParameter = relativeDeathRateParameter;
        addParameter(relativeDeathRateParameter);
        relativeDeathRateParameter.addBounds(new Parameter.DefaultBounds(1.0, 0.0, 1));

        this.sampleProbability = sampleProbability;
        if( sampleProbability != null ) {
            addParameter(sampleProbability);
            sampleProbability.addBounds(new Parameter.DefaultBounds(1.0, 0.0, 1));
        }
    }

    public double getR() {
        return birthDiffRateParameter.getParameterValue(0);
    }

    public double getA() {
        return relativeDeathRateParameter.getParameterValue(0);
    }

    public double getRho() {
        return sampleProbability != null ? sampleProbability.getParameterValue(0) : 1.0;
    }

    public double logTreeProbability(int taxonCount) {
        return logGamma(taxonCount + 1) +
                (taxonCount - 1) * Math.log(getR() * getRho()) + taxonCount * Math.log(1 - getA());
    }

    public double logNodeProbability(Tree tree, NodeRef node) {
        final double height = tree.getNodeHeight(node);
        final double rho = getRho();
        final double mrh = -getR() * height;
        final double z = Math.log(rho + ((1-rho) - getA()) * Math.exp(mrh));
        double l = -2 * z + mrh;

        if (tree.getRoot() == node) {
            l += mrh - z;
        }
        return l;
    }

    public boolean includeExternalNodesInLikelihoodCalculation() {
        return false;
    }
}
@@ -91,16 +91,61 @@
 e.Tree;%0A
+import dr.evomodelxml.BirthDeathModelParser;%0A
 import d
@@ -1025,10 +1025,9 @@
 L = 
-%22b
+B
 irth
@@ -1036,17 +1036,40 @@
 athModel
-%22
+Parser.BIRTH_DEATH_MODEL
 ;%0A%0A    P

avoid string duplications