package beast.math.distributions;


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

import beast.core.*;
import org.apache.commons.math.distribution.GammaDistribution;
import org.apache.commons.math.distribution.GammaDistributionImpl;

import beast.core.Input.Validate;
import beast.core.parameter.RealParameter;
import beast.math.distributions.LogNormalDistributionModel.LogNormalImpl;



/**
 * 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 useLogNormal is set, a log normal distribution is used instead of a Gamma. " +
        "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 {

    final public Input<Boolean> isJeffreysInput = new Input<>("jeffreys", "use Jeffrey's prior (default false)", false);
    final public Input<Boolean> isReverseInput = new Input<>("reverse", "parameter in reverse (default false)", false);
    final public Input<Boolean> useLogInput = new Input<>("uselog", "use logarithm of parameter values (default false)", false);
    final public Input<Double> shapeInput = new Input<>("shape", "shape parameter of the Gamma distribution (default 1.0 = exponential distribution) " +
    		" or precision parameter if the log normal is used.", 1.0);
    final public Input<Function> parameterInput = new Input<>("parameter", "chain parameter to calculate distribution over", Validate.REQUIRED);
    final public Input<Boolean> useLogNormalInput = new Input<>("useLogNormal", "use Log Normal distribution instead of Gamma (default false)", false);

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

    @Override
    public void initAndValidate() {
        reverse = isReverseInput.get();
        jeffreys = isJeffreysInput.get();
        uselog = useLogInput.get();
        shape = shapeInput.get();
        chainParameter = parameterInput.get();
        useLogNormal = useLogNormalInput.get();
        gamma = new GammaDistributionImpl(shape, 1);
        logNormal = new LogNormalDistributionModel().new LogNormalImpl(1, 1);
    }


    /**
     * Get the log likelihood.
     *
     * @return the log likelihood.
     */
    @SuppressWarnings("deprecation")
	@Override
    public double calculateLogP() {
        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);

            if (useLogNormal) {
	            final double sigma = 1.0 / shape; // shape = precision
	            // convert mean to log space
	            final double M = Math.log(mean) - (0.5 * sigma * sigma);
	            logNormal.setMeanAndStdDev(M, sigma);
	            logP += logNormal.logDensity(x);
            } else {
                final double scale = mean / shape;
                gamma.setBeta(scale);
                logP += gamma.logDensity(x);
            }
        }
        return logP;
    }

    private double getChainValue(int i) {
        if (uselog) {
            return Math.log(chainParameter.getArrayValue(index(i)));
        } else {
            return chainParameter.getArrayValue(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) {
    }
}


@@ -82,29 +82,8 @@
 m;%0A%0A
-import beast.core.*;%0A
 impo
@@ -209,16 +209,104 @@
 nImpl;%0A%0A
+import beast.core.Description;%0Aimport beast.core.Distribution;%0Aimport beast.core.Input;%0A
 import b
@@ -331,16 +331,41 @@
 lidate;%0A
+import beast.core.State;%0A
 import b
@@ -1727,24 +1727,29 @@
 c Input%3C
-Function
+RealParameter
 %3E parame
@@ -2197,16 +2197,21 @@
 ate 
-Function
+RealParameter
  cha
@@ -3983,37 +3983,32 @@
 ainParameter.get
-Array
 Value(index(i)))
@@ -4067,13 +4067,8 @@
 .get
-Array
 Valu

Revert "M.C. distrib now takes a Function as input."