package beast.core.util;

import java.io.PrintStream;
import java.util.ArrayList;
import java.util.List;

import beast.core.BEASTObject;
import beast.core.CalculationNode;
import beast.core.Description;
import beast.core.Function;
import beast.core.Input;
import beast.core.Input.Validate;
import beast.core.Loggable;
import beast.core.parameter.BooleanParameter;
import beast.core.parameter.IntegerParameter;



@Description("calculates sum of a valuable")
public class Sum extends CalculationNode implements Function, Loggable {
    final public Input<List<Function>> functionInput = new Input<>("arg", "argument to be summed", new ArrayList<>(), Validate.REQUIRED);

    enum Mode {integer_mode, double_mode}

    Mode mode;

    boolean needsRecompute = true;
    double sum = 0;
    double storedSum = 0;

    @Override
    public void initAndValidate() {
        List<Function> valuable = functionInput.get();
        mode = Mode.integer_mode;
        for (Function v : valuable) {
	        if (!(v instanceof IntegerParameter || v instanceof BooleanParameter)) {
	            mode = Mode.double_mode;
	        }
        }
    }

    @Override
    public int getDimension() {
        return 1;
    }

    @Override
    public double getArrayValue() {
        if (needsRecompute) {
            compute();
        }
        return sum;
    }

    /**
     * do the actual work, and reset flag *
     */
    void compute() {
        sum = 0;
        for (Function v : functionInput.get()) {
	        for (int i = 0; i < v.getDimension(); i++) {
	            sum += v.getArrayValue(i);
	        }
        }
        needsRecompute = false;
    }

    @Override
    public double getArrayValue(int dim) {
        if (dim == 0) {
            return getArrayValue();
        }
        return Double.NaN;
    }

    /**
     * CalculationNode methods *
     */
    @Override
    public void store() {
        storedSum = sum;
        super.store();
    }

    @Override
    public void restore() {
        sum = storedSum;
        super.restore();
    }

    @Override
    public boolean requiresRecalculation() {
        needsRecompute = true;
        return true;
    }

    /**
     * Loggable interface implementation follows
     */
    @Override
    public void init(PrintStream out) {
        out.print("sum(" + ((BEASTObject) functionInput.get().get(0)).getID() + ")\t");
    }

    @Override
    public void log(long sampleNr, PrintStream out) {
        double sum = 0;
        for (Function v : functionInput.get()) {
	        for (int i = 0; i < v.getDimension(); i++) {
	            sum += v.getArrayValue(i);
	        }
        }
        if (mode == Mode.integer_mode) {
            out.print((int) sum + "\t");
        } else {
            out.print(sum + "\t");
        }
    }

    @Override
    public void close(PrintStream out) {
        // nothing to do
    }

} // class Sum

@@ -403,17 +403,50 @@
 ameter;%0A
+import beast.evolution.tree.Tree;
 %0A
-
 %0A%0A@Descr
@@ -696,16 +696,407 @@
 IRED);%0A%0A
+    final public Input%3CTree%3E treeInput = new Input%3C%3E(%22tree%22, %22the tree corresponding to the function to be summed, indexing by node numbers assumed.%22, Validate.OPTIONAL);%0A%0A    final public Input%3CBoolean%3E ignoreZeroBranchLengthsInput = new Input%3C%3E(%22ignoreZeroBranchLengths%22, %22true if quantities in the argument should only be summed for non-zero branch lengths.%22, false, Validate.OPTIONAL);%0A%0A
     enum
@@ -1228,16 +1228,77 @@
 m = 0;%0A%0A
+    Tree tree;%0A    boolean ignoreZeroBranchLengths = false;%0A%0A
     @Ove
@@ -1599,32 +1599,134 @@
     %7D%0A        %7D%0A
+        tree = treeInput.get();%0A        ignoreZeroBranchLengths = ignoreZeroBranchLengthsInput.get();%0A
     %7D%0A%0A    @Over
@@ -2017,32 +2017,91 @@
        sum = 0;%0A
+        if (tree != null && ignoreZeroBranchLengths) %7B%0A    
         for (Fun
@@ -2125,33 +2125,40 @@
 nInput.get()) %7B%0A
-%09
+        
         for (int
@@ -2198,49 +2198,377 @@
 ) %7B%0A
-%09            sum += v.getArrayValue(i);%0A%09
+                    if (!tree.getNode(i).isDirectAncestor()) %7B%0A                        sum += v.getArrayValue(i);%0A                    %7D%0A                %7D%0A            %7D%0A        %7D else %7B%0A%0A            for (Function v : functionInput.get()) %7B%0A                for (int i = 0; i %3C v.getDimension(); i++) %7B%0A                    sum += v.getArrayValue(i);%0A                %7D%0A    
     

Add sampled-ancestor aware option to the Sum Function.