/*
 * ExchangeOperator.java
 *
 * Copyright (C) 2002-2006 Alexei Drummond and Andrew Rambaut
 *
 * This file is part of BEAST.
 * See the NOTICE file distributed with this work for additional
 * information regarding copyright ownership and licensing.
 *
 * BEAST is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as
 * published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 *
 *  BEAST is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with BEAST; if not, write to the
 * Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
 * Boston, MA  02110-1301  USA
 */

package dr.evomodel.operators;

import dr.evolution.tree.MutableTree;
import dr.evolution.tree.NodeRef;
import dr.evomodel.tree.TreeModel;
import dr.inference.operators.OperatorFailedException;
import dr.math.MathUtils;
import dr.xml.*;

/**
 *  Uniform height change of several related tree nodes at once.
 *
 *  Currently supported - a parent-son pair and parent-two sons triplet.
 *
 *  @author  Joseph Heled
 */
public class TreeUniform extends AbstractTreeOperator {

    public static final String TREE_UNIFORM = "treeUniform";

    private final int nodesToMove;

    //private static final int MAX_TRIES = 100;
    private final TreeModel tree;

    public TreeUniform(int n, TreeModel tree, double weight) {
        assert n == 2 || n == 3;
        
        this.tree = tree;
        this.nodesToMove = n;
        setWeight(weight);
    }

    public double doOperation() throws OperatorFailedException {
        tree.beginTreeEdit();

        switch( nodesToMove) {
            case 2: move2(); break;
            case 3: move3(); break;
        }

        try {
            tree.endTreeEdit();
        } catch( MutableTree.InvalidTreeException ite ) {
            throw new OperatorFailedException(ite.toString());
        }

        return 0;
    }


    // insure two elements are sorted in reverse order
    private void sorted(double[] h, int i, int j) {
        if( h[i] < h[j] ) {
           exch(h, i, j);
        }
    }

    private void exch(double[] h, int i, int j) {
        final double t = h[i];
        h[i] = h[j];
        h[j] = t;
    }

    private void move3() throws OperatorFailedException {
        final int nInternalNodes = tree.getInternalNodeCount();

        final NodeRef root = tree.getRoot();

        //  find an internal node with 2 internal children
        NodeRef i = root;
        int nTries = 0;

        while( root == i || tree.isExternal(tree.getChild(i, 0)) || tree.isExternal(tree.getChild(i, 1)) ) {
            i = tree.getInternalNode(MathUtils.nextInt(nInternalNodes));
            if( ++nTries > 100 ) {
                throw new OperatorFailedException("no node found");
            }
        }

        final NodeRef iParent = tree.getParent(i);

        final NodeRef child0 = tree.getChild(i, 0);
        final NodeRef child1 = tree.getChild(i, 1);

        final double hMin0 = getMaxChildHeight(child0);
        final double hMin1 = getMaxChildHeight(child1);
        // upper limit for i
        final double hMax = tree.getNodeHeight(iParent);

        final double mx = Math.max(hMin0, hMin1);
        final double mn = Math.min(hMin0, hMin1);

        final double d1 = hMax - mx;

        // factor of d1^2 dropped
        final double
                a0 = Math.abs(hMin0 - hMin1) / 2,
                a1 = d1 / 3;

        final double th = a0 / (a0 + a1);

        double[] h = new double[3];

        final int il = (hMin0 < hMin1) ? 0 : 1;

        for(int k = 0; k < 3; ++k) {
            h[k] = MathUtils.uniform(mx, hMax);
        }

        if( MathUtils.nextDouble() < th ) {
            h[1+il] = MathUtils.uniform(mn,mx);
            sorted(h, 0, 2-il);
        } else {
            int iMax =  h[0] > h[1] ? 0 : 1;
            iMax = h[iMax] < h[2] ? 2 : iMax;
            exch(h, 0, iMax);
        }

        assert hMax > h[0] && h[0] > h[1] && h[0] > h[2] && h[1] > hMin0 && h[2] > hMin1;
        
        NodeRef[] nodes = {i, child0, child1};
        for(int k = 0; k < nodes.length; ++k) {
            tree.setNodeHeight(nodes[k], h[k]);
            tree.pushTreeChangedEvent(nodes[k]);
        }
    }

    private void move2() {
        final int nInternalNodes = tree.getInternalNodeCount();
        assert nInternalNodes > 3;

        final NodeRef root = tree.getRoot();

        NodeRef i = root;

        while( root == i || tree.getParent(i) == root ) {
            i = tree.getInternalNode(MathUtils.nextInt(nInternalNodes)) ;
        }

        final NodeRef iParent = tree.getParent(i);
        final NodeRef iGrandParent = tree.getParent(iParent);

        // lower limit for node (max height of children)
        final double hMin = getMaxChildHeight(i);
        // upper limit for parent (it's parent height)
        final double hMax = tree.getNodeHeight(iGrandParent);
        // lower bound for parent
        final double hLim0 = Math.max(hMin, tree.getNodeHeight(getOtherChild(iParent, i)));

        final double d1 = hMax - hLim0;
        // factor of d1 dropped
        final double a0 = (hLim0 - hMin),
                     a1 = d1/2;

        final double th = a0 / (a0 + a1);
        double[] h = new double[2];

        if( MathUtils.nextDouble() < th ) {
           h[0] = MathUtils.uniform(hMin, hLim0);
           h[1] = MathUtils.uniform(hLim0, hMax);
        } else {
            for(int k = 0; k < 2; ++k) {
              h[k] = MathUtils.uniform(hLim0, hMax);
            }
            // sorted(h, 1, 0) ??
            if( h[0] > h[1] ) {
                double t = h[0];
                h[0] = h[1];
                h[1] = t;
            }
        }

        tree.setNodeHeight(iParent, h[1]);
        tree.setNodeHeight(i, h[0]);

        tree.pushTreeChangedEvent(iParent);
        tree.pushTreeChangedEvent(i);
    }

    private NodeRef getOtherChild(NodeRef n, NodeRef c) {
        for(int nc = 0; nc < tree.getChildCount(n); ++nc ) {
            final NodeRef child = tree.getChild(n, nc);
            if( child != c ) {
               return child;
           }
        }
        assert true;
        return null;
    }

    private double getMaxChildHeight(NodeRef n) {
        double hMax = -1;
        for(int nc = 0; nc < tree.getChildCount(n); ++nc ) {
            hMax = Math.max(hMax, tree.getNodeHeight(tree.getChild(n, nc) ));
        }
        return hMax;
    }


    public String getPerformanceSuggestion() {
        return "";
    }

    public String getOperatorName() {
        return "treeUniform(" + nodesToMove + "," + tree.getId() + ")";
    }


    public static XMLObjectParser PARSSR = new AbstractXMLObjectParser() {

        public String getParserName() {
            return TREE_UNIFORM;
        }

        public Object parseXMLObject(XMLObject xo) throws XMLParseException {

            TreeModel treeModel = (TreeModel) xo.getChild(TreeModel.class);
            final double weight = xo.getDoubleAttribute("weight");
            final int n = xo.getAttribute("count", 2);
            if( ! ( n == 2 || n == 3) ) {
              throw new XMLParseException("Sorry, only moves of 2 or 3 nodes implemented.");
            }
            return new TreeUniform(n, treeModel, weight);
        }

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

        public String getParserDescription() {
            return "Simultanouesly change height of two nodes.";
        }

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

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

        private final XMLSyntaxRule[] rules = {
                AttributeRule.newDoubleRule("weight"),
                AttributeRule.newIntegerRule("count", true),
                new ElementRule(TreeModel.class)};

    };
}