package beast.math.distributions;

import java.util.List;
import java.util.Random;

import org.apache.commons.math.distribution.GammaDistribution;
import org.apache.commons.math.distribution.GammaDistributionImpl;

import beast.core.Description;
import beast.core.Distribution;
import beast.core.Input;
import beast.core.Input.Validate;
import beast.core.State;
import beast.core.parameter.RealParameter;

/**
 * Initial version Ported from Beast 1.7 ExponentialMarkovModel
 */
@Description("A class that produces a distribution chaining values in a parameter through the Gamma distribution. " +
        "The value of a parameter is assumed to be Gamma distributed with mean as the previous value in the parameter. " +
        "If a Jeffrey's prior is used, the first value is assumed to be distributed as 1/x, otherwise it is assumed to be uniform. " +
        "Handy for population parameters. ")
public class MarkovChainDistribution extends Distribution {

    public Input<Boolean> bIsJeffreysInput = new Input<Boolean>("jeffreys", "use Jeffrey's prior (default false)", false);
    public Input<Boolean> bIsReverseInput = new Input<Boolean>("reverse", "parameter in reverse (default false)", false);
    public Input<Boolean> bUseLogInput = new Input<Boolean>("uselog", "use logarithm of parameter values (default false)", false);
    public Input<Double> fShapeInput = new Input<Double>("shape", "shape parameter of the Gamma distribution (default 1.0 = exponential distribution)", 1.0);
    public Input<RealParameter> parameterInput = new Input<RealParameter>("parameter", "chain parameter to calculate distribution over", Validate.REQUIRED);

    // **************************************************************
    // Private instance variables
    // **************************************************************
    private RealParameter chainParameter = null;
    private boolean jeffreys = false;
    private boolean reverse = false;
    private boolean uselog = false;
    private double shape = 1.0;
    GammaDistribution gamma;

    @Override
    public void initAndValidate() throws Exception {
        reverse = bIsReverseInput.get();
        jeffreys = bIsJeffreysInput.get();
        uselog = bUseLogInput.get();
        shape = fShapeInput.get();
        chainParameter = parameterInput.get();
        gamma = new GammaDistributionImpl(shape, 1);
    }


    /**
     * Get the log likelihood.
     *
     * @return the log likelihood.
     */
    @Override
    public double calculateLogP() throws Exception {
        logP = 0.0;
        // jeffreys Prior!
        if (jeffreys) {
            logP += -Math.log(getChainValue(0));
        }
        for (int i = 1; i < chainParameter.getDimension(); i++) {
            final double mean = getChainValue(i - 1);
            final double x = getChainValue(i);
            //logL += dr.math.distributions.ExponentialDistribution.logPdf(x, 1.0/mean);

            final double scale = mean / shape;
            gamma.setBeta(scale);
            logP += gamma.logDensity(x);//logPdf(x, shape, scale);
        }
        return logP;
    }

    private double getChainValue(int i) {
        if (uselog) {
            return Math.log(chainParameter.getValue(index(i)));
        } else {
            return chainParameter.getValue(index(i));
        }
    }

    private int index(int i) {
        if (reverse)
            return chainParameter.getDimension() - i - 1;
        else
            return i;
    }

    @Override
    public List<String> getArguments() {
        return null;
    }

    @Override
    public List<String> getConditions() {
        return null;
    }

    @Override
    public void sample(State state, Random random) {
    }
}


@@ -28,16 +28,17 @@
 tions;%0A%0A
+%0A
 import j
@@ -398,16 +398,90 @@
 rameter;
+%0Aimport beast.math.distributions.LogNormalDistributionModel.LogNormalImpl;
 %0A%0A/**%0A *
@@ -783,24 +783,109 @@
 ameter. %22 +%0A
+%09%09%22If useLogNormal is set, a log normal distribution is used instead of a Gamma. %22 +%0A
         %22If 
@@ -1638,16 +1638,78 @@
 ibution)
+ %22 +%0A    %09%09%22 or precision parameter if the log normal is used.
 %22, 1.0);
@@ -1865,16 +1865,169 @@
 QUIRED);
+%0A    public Input%3CBoolean%3E useLogNormalInput = new Input%3CBoolean%3E(%22useLogNormal%22, %22use Log Normal distribution instead of Gamma (default false)%22, false);
 %0A%0A    //
@@ -2414,16 +2414,71 @@
 n gamma;
+%0A    LogNormalImpl logNormal;%0A    boolean useLogNormal;
 %0A%0A    @O
@@ -2733,32 +2733,80 @@
 terInput.get();%0A
+        useLogNormal = useLogNormalInput.get();%0A
         gamma = 
@@ -2838,24 +2838,102 @@
 (shape, 1);%0A
+        logNormal = new LogNormalDistributionModel().new LogNormalImpl(1, 1);%0A
     %7D%0A%0A%0A    
@@ -3381,16 +3381,17 @@
 lue(i);%0A
+%0A
         
@@ -3398,86 +3398,329 @@
     
-//logL += dr.math.distributions.ExponentialDistribution.logPdf(x, 1.0/mean);%0A%0A
+if (useLogNormal) %7B%0A%09            final double sigma = 1.0 / shape; // shape = precision%0A%09            // convert mean to log space%0A%09            final double M = Math.log(mean) - (0.5 * sigma * sigma);%0A%09            logNormal.setMeanAndStdDev(M, sigma);%0A%09            logP += logNormal.logDensity(x);%0A            %7D else %7B%0A   
     
@@ -3719,32 +3719,33 @@
 %0A               
+ 
 final double sca
@@ -3771,24 +3771,28 @@
             
+    
 gamma.setBet
@@ -3805,32 +3805,36 @@
 e);%0A            
+ 
+   
 logP += gamma.lo
@@ -3849,34 +3849,22 @@
 (x);
-//logPdf(x, shape, scale);
+%0A            %7D
 %0A   

add lognormal flag to Markov chain distribution