package dr.inference.model;
import dr.evolution.tree.NodeRef;
import dr.evolution.util.Taxon;
import dr.evomodel.tree.TreeModel;
import dr.xml.*;

/**
 * @author Alexander Fisher
 */
public class MaskFromTree extends Parameter.Abstract implements ModelListener {

    public static final String MASK_FROM_TREE = "maskFromTree";
    private TreeModel tree;
    private Taxon taxon;
    private boolean ancestralMaskBranchKnown = false;
    private Parameter maskParameter;

    public MaskFromTree(TreeModel tree, Taxon referenceTaxon) {
        this.tree = tree;
        this.taxon = referenceTaxon;
        int numBranches = tree.getNodeCount() - 1;
        this.maskParameter = new Parameter.Default(numBranches, 1.0);
        tree.addModelListener(this);
        tree.addModelRestoreListener(this);
        updateMask();
    }

    void updateMask() {
       int nodeNumber =  tree.getTaxonIndex(taxon.getId());
       NodeRef node = tree.getNode(nodeNumber);
       NodeRef root = tree.getRoot();

       while(tree.getParent(node) != root){
           node = tree.getParent(node);
       }

        int maskIndex = node.getNumber();
        maskParameter.setParameterValue(maskIndex, 0.0);
        ancestralMaskBranchKnown = true;
    }


    public int getDimension(){
        return maskParameter.getDimension();
    }

    @Override
    public double getParameterValue(int dim) {
        if (!ancestralMaskBranchKnown) { updateMask(); }
        return maskParameter.getParameterValue(dim);
    }

    @Override
    public void setParameterValue(int dim, double value) {
        throw new RuntimeException("Not yet implemented");
    }

    @Override
    public void setParameterValueQuietly(int dim, double value) {
        throw new RuntimeException("Not yet implemented");
    }

    @Override
    public void setParameterValueNotifyChangedAll(int dim, double value) {
        throw new RuntimeException("Not yet implemented");
    }

    @Override
    public String getParameterName() {
        return null;
    }

    @Override
    public void addBounds(Bounds<Double> bounds) {
        throw new RuntimeException("Not yet implemented");
    }

    @Override
    public Bounds<Double> getBounds() {
        return null;
    }

    @Override
    public void addDimension(int index, double value) {
        throw new RuntimeException("Not yet implemented");
    }

    @Override
    public double removeDimension(int index) {
        return 0;
    }

    @Override
    protected void storeValues() {
        throw new RuntimeException("Not yet implemented");
    }

    @Override
    protected void restoreValues() {
        throw new RuntimeException("Not yet implemented");
    }

    @Override
    protected void acceptValues() {
        throw new RuntimeException("Not yet implemented");
    }

    @Override
    protected void adoptValues(Parameter source) {
        throw new RuntimeException("Not yet implemented");
    }

    @Override
    public void modelChangedEvent(Model model, Object object, int index) {
        ancestralMaskBranchKnown=false;
    }

    @Override
    public void modelRestored(Model model) {
        ancestralMaskBranchKnown=false;
    }


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

    public static XMLObjectParser PARSER = new AbstractXMLObjectParser() {

        public String getParserName() {
            return MASK_FROM_TREE;
        }

        public Object parseXMLObject(XMLObject xo) throws XMLParseException {

            TreeModel tree = (TreeModel) xo.getChild(TreeModel.class);

            Taxon referenceTaxon = (Taxon) xo.getChild(Taxon.class);


            MaskFromTree maskFromTree = new MaskFromTree(tree, referenceTaxon);
            return maskFromTree;
        }

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

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

        private XMLSyntaxRule[] rules = new XMLSyntaxRule[] {
                new ElementRule(TreeModel.class),
                new ElementRule(Taxon.class)
        };

        public String getParserDescription() {
            return "Masks ancestral (off-root) branch to a specific reference taxon on a random tree";
        }

        public Class getReturnType() { return MaskFromTree.class; }
    };
}

@@ -17,24 +17,25 @@
 ence.model;%0A
+%0A
 import dr.ev
@@ -466,16 +466,74 @@
 ameter;%0A
+// TODO: remove maskParameter, use store/restore machinery
 %0A    pub
@@ -910,24 +910,25 @@
 () %7B%0A       
+ 
 int nodeNumb
@@ -931,17 +931,16 @@
 Number =
- 
  tree.ge
@@ -967,16 +967,17 @@
 tId());%0A
+ 
        N
@@ -1019,24 +1019,25 @@
 er);%0A       
+ 
 NodeRef root
@@ -1052,24 +1052,25 @@
 getRoot();%0A%0A
+ 
        while
@@ -1069,16 +1069,17 @@
    while
+ 
 (tree.ge
@@ -1100,18 +1100,20 @@
 != root)
+ 
 %7B%0A
+ 
         
@@ -1144,16 +1144,17 @@
 (node);%0A
+ 
        %7D
@@ -1332,16 +1332,17 @@
 ension()
+ 
 %7B%0A      
@@ -1488,16 +1488,28 @@
 Known) %7B
+%0A           
  updateM
@@ -1514,16 +1514,24 @@
 eMask();
+%0A       
  %7D%0A     
@@ -3133,33 +3133,35 @@
 lMaskBranchKnown
-=
+ = 
 false;%0A    %7D%0A%0A  
@@ -3249,17 +3249,19 @@
 nchKnown
-=
+ = 
 false;%0A 
@@ -4201,24 +4201,36 @@
 taxRules() %7B
+%0A           
  return rule
@@ -4231,16 +4231,24 @@
 n rules;
+%0A       
  %7D%0A%0A    
@@ -4302,17 +4302,16 @@
 axRule%5B%5D
- 
 %7B%0A      
@@ -4610,16 +4610,28 @@
 Type() %7B
+%0A           
  return 
@@ -4649,16 +4649,24 @@
 e.class;
+%0A       
  %7D%0A    %7D

format code, add todo