package dr.evomodel.branchratemodel.shrinkage;


import dr.evolution.tree.TreeTrait;
import dr.evolution.tree.TreeTraitProvider;
import dr.evomodel.branchratemodel.AutoCorrelatedBranchRatesDistribution;
import dr.inference.distribution.ParametricMultivariateDistributionModel;
import dr.inference.model.Parameter;
import dr.xml.*;

/**
 * @author Alexander Fisher
 */
public class IncrementClassifier implements TreeTraitProvider {
// todo: find magic default epsilon from prior simulation
    public static final String PARAMETER_CLASSIFIER = "parameterClassifier";
    public static final String EPSILON = "epsilon";
    public static final String TARGET_PROBABILITY = "targetProbability";

    private AutoCorrelatedBranchRatesDistribution brm;
    private double epsilon;
    private boolean epsilonKnown = false;
    private double targetProb;
    private int dim;
    private ParametricMultivariateDistributionModel distribution;

    private double[] classified;

    public IncrementClassifier(AutoCorrelatedBranchRatesDistribution brm, double epsilon, double targetProb) {
        this.brm = brm;
        this.distribution = brm.getPrior();

        if (epsilon != 0.0) {
            this.epsilon = epsilon;
            epsilonKnown = true;
        } else if (targetProb != 0.0) {
            this.targetProb = targetProb;
            this.epsilon = findEpsilonFromTargetProb(targetProb);
        }

        this.dim = brm.getDimension();
        this.classified = new double[dim];
        classify();
    }

    private double findEpsilonFromTargetProb(double targetProb) {
        // finds eps such that p(abs(increment)) < epsilon = targetProb
        throw new RuntimeException("Target probabiltiy not yet implemented.");
    }

    private void classify() {
        double increment;
        if(!epsilonKnown){
            findEpsilonFromTargetProb(targetProb);
        }
        for (int i = 0; i < dim; i++) {
            increment = brm.getIncrement(i);
            classified[i] = 1; // is an increment
            if (increment < epsilon) {
                classified[i] = 0; // is not an increment
            }
        }
    }

    @Override
    public TreeTrait[] getTreeTraits() {
        return new TreeTrait[0];
    }

    @Override
    public TreeTrait getTreeTrait(String key) {
        return null;
    }


    // **************************************************************
    // XMLObjectParser
    // **************************************************************

    public static XMLObjectParser PARSER = new AbstractXMLObjectParser() {

        public String getParserName() {
            return PARAMETER_CLASSIFIER;
        }

        public Object parseXMLObject(XMLObject xo) throws XMLParseException {

//            Parameter parameter = (Parameter) xo.getChild(Parameter.class);
            AutoCorrelatedBranchRatesDistribution brm = (AutoCorrelatedBranchRatesDistribution) xo.getChild(AutoCorrelatedBranchRatesDistribution.class);

            if (xo.hasAttribute(EPSILON) && xo.hasAttribute(TARGET_PROBABILITY)) {
                throw new XMLParseException("Cannot set both epsilon and target probability.");
            }

            double epsilon = xo.getAttribute(EPSILON, 0);
            double targetProbability = xo.getAttribute(TARGET_PROBABILITY, 0);

            if (epsilon < 0) {
                throw new XMLParseException("epsilon must be positive.");
            }
            else if (targetProbability < 0) {
                throw new XMLParseException("target probability must be positive.");
            }

            IncrementClassifier incrementClassifier = new IncrementClassifier(brm, epsilon, targetProbability);
            return incrementClassifier;
        }

        //************************************************************************
        // AbstractXMLObjectParser implementation
        //************************************************************************

        public XMLSyntaxRule[] getSyntaxRules() {
            return rules;
        }

        @Override
        public String getParserDescription() {
            return "Reports 0 or 1 depending on whether a model parameter falls within a specified density range.";
        }

        @Override
        public Class getReturnType() {
            return IncrementClassifier.class;
        }

        private XMLSyntaxRule[] rules = new XMLSyntaxRule[]{
                new ElementRule(Parameter.class),
                AttributeRule.newDoubleRule(EPSILON, true),
                AttributeRule.newDoubleRule(TARGET_PROBABILITY, true),
        };
    };
}

@@ -274,43 +274,174 @@
 el;%0A
-import dr.inference.model.Parameter
+%0Aimport dr.inference.distribution.shrinkage.JointBayesianBridgeDistributionModel;%0Aimport dr.inference.model.Parameter;%0Aimport dr.math.distributions.NormalDistribution
 ;%0Aim
@@ -584,23 +584,20 @@
 default 
-epsilon
+s.d.
  from pr
@@ -606,24 +606,58 @@
 r simulation
+ (then use targetProb to find eps)
 %0A    public 
@@ -976,94 +976,290 @@
 lse;
-%0A    private double targetProb;%0A    private int dim;%0A    private ParametricMultivariat
+ // todo: cleanup -- probably don't need to dynamically update epsilon%0A    private double targetProb;%0A    private int dim;%0A%0A    private double sd = 1.0; //todo: put magic number here.%0A//    private ParametricMultivariateDistributionModel distribution;%0A    private JointBayesianBridg
 eDis
@@ -1489,22 +1489,147 @@
 n = 
-brm.getPrior()
+(distribution instanceof JointBayesianBridgeDistributionModel) ?%0A                (JointBayesianBridgeDistributionModel) distribution : null
 ;%0A%0A 
@@ -2095,10 +2095,13 @@
 hat 
-p(
+Prob%7B
 abs(
@@ -2114,19 +2114,15 @@
 ent)
-)
  %3C eps
-ilon
+%7D
  = t
@@ -2143,76 +2143,199 @@
     
-throw new RuntimeException(%22T
+double probToFindEpsilon = (t
 arget
- probabiltiy not yet implemented.%22
+Prob / 2) + 0.5;%0A        NormalDistribution norm = new NormalDistribution(0, sd);%0A        epsilonKnown = true;%0A        return norm.quantile(probToFindEpsilon
 );%0A 
@@ -2634,16 +2634,25 @@
      if 
+(Math.abs
 (increme
@@ -2653,16 +2653,17 @@
 ncrement
+)
  %3C epsil

tweak plan