package dr.evomodel.treedatalikelihood.discrete;

import dr.evolution.tree.NodeRef;
import dr.evolution.tree.Tree;
import dr.evomodel.treedatalikelihood.BeagleDataLikelihoodDelegate;
import dr.evomodel.treedatalikelihood.TreeDataLikelihood;
import dr.inference.hmc.GradientWrtParameterProvider;
import dr.inference.loggers.Loggable;
import dr.inference.model.Parameter;
import dr.xml.Reportable;

public class LocalBranchRateGradientForDiscreteTrait extends DiscreteTraitBranchRateGradient
        implements GradientWrtParameterProvider, Reportable, Loggable {

    public LocalBranchRateGradientForDiscreteTrait(String traitName,
                                                   TreeDataLikelihood treeDataLikelihood,
                                                   BeagleDataLikelihoodDelegate likelihoodDelegate,
                                                   Parameter rateParameter,
                                                   boolean useHessian) {
        super(traitName, treeDataLikelihood, likelihoodDelegate, rateParameter, useHessian);
    }

    @Override
    public double[] getGradientLogDensity() {
        double[] nodeGradient = super.getGradientLogDensity();
        double[] result = new double[tree.getNodeCount() - 1];

        for (int i = 0; i < tree.getNodeCount(); i++) {
            NodeRef node = tree.getNode(i);
            if(!tree.isRoot(node)){
                final int parameterIndex = branchRateModel.getParameterIndexFromNode(node);
                result[parameterIndex] = getSubTreeGradient(tree, node, nodeGradient);
            }
        }
        return result;
    }

    private double getSubTreeGradient(Tree tree, NodeRef node, final double[] nodeGradient) {
        final double chainGradient =  branchRateModel.getBranchRateDifferential(tree, node);
        final double branchGradient = nodeGradient[branchRateModel.getParameterIndexFromNode(node)];
        if (tree.isExternal(node)) {
            return chainGradient * branchGradient;
        } else {
            double sum = chainGradient * branchGradient;
            for (int i = 0; i < tree.getChildCount(node); i++) {
                sum += getSubTreeGradient(tree, tree.getChild(node, i), nodeGradient);
            }
            return sum;
        }
    }
}

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

add author