<commit_before>/*
 * TopologyTracer.java
 *
 * Copyright (c) 2002-2017 <NAME>, <NAME> and <NAME>
 *
 * 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.evolution.tree;

import dr.evolution.io.Importer;
import dr.evolution.io.NewickImporter;
import dr.util.Pair;

import java.io.IOException;
import java.util.*;

/**
 * @author <NAME>
 * Path difference metric according to Kendall & Colijn (2015)
 */
public class KCPathDifferenceMetric {

    private Tree focalTree;
    private int dim, externalNodeCount;
    private double[] focalSmallM, focalLargeM;

    public KCPathDifferenceMetric() {

    }

    public KCPathDifferenceMetric(Tree focalTree) {
        this.focalTree = focalTree;
        this.externalNodeCount = focalTree.getExternalNodeCount();
        this.dim = (externalNodeCount-2)*(externalNodeCount-1)+externalNodeCount;
        this.focalSmallM = new double[dim];
        this.focalLargeM = new double[dim];

        processTree(focalTree, focalSmallM, focalLargeM, dim);
    }

    public ArrayList<Double> getMetric_old(Tree tree, ArrayList<Double> lambda) {

        //check if taxon lists are in the same order!!
        if (focalTree.getExternalNodeCount() != tree.getExternalNodeCount()) {
            throw new RuntimeException("Different number of taxa in both trees.");
        } else {
            for (int i = 0; i < focalTree.getExternalNodeCount(); i++) {
                if (!focalTree.getNodeTaxon(focalTree.getExternalNode(i)).getId().equals(tree.getNodeTaxon(tree.getExternalNode(i)).getId())) {
                    throw new RuntimeException("Mismatch between taxa in both trees: " + focalTree.getNodeTaxon(focalTree.getExternalNode(i)).getId() + " vs. " + tree.getNodeTaxon(tree.getExternalNode(i)).getId());
                }
            }
        }

        double[] smallMTwo = new double[dim];
        double[] largeMTwo = new double[dim];

        int index = 0;
        for (int i = 0; i < tree.getExternalNodeCount(); i++) {
            for (int j = i+1; j < tree.getExternalNodeCount(); j++) {
                //get two leaf nodes
                NodeRef nodeOne = tree.getExternalNode(i);
                NodeRef nodeTwo = tree.getExternalNode(j);

                //get common ancestor of 2 leaf nodes
                NodeRef MRCA = TreeUtils.getCommonAncestor(tree, nodeOne, nodeTwo);

                int edges = 0;
                double branchLengths = 0.0;
                while (MRCA != tree.getRoot()) {
                    edges++;
                    branchLengths += tree.getNodeHeight(tree.getParent(MRCA)) - tree.getNodeHeight(MRCA);
                    MRCA = tree.getParent(MRCA);
                }
                smallMTwo[index] = edges;
                largeMTwo[index] = branchLengths;
                index++;
            }
        }

        //fill out arrays further
        index = 0;
        for (int i = (tree.getExternalNodeCount()-1)*(tree.getExternalNodeCount()-2); i < dim; i++) {
            smallMTwo[i] = 1.0;
            largeMTwo[i] = tree.getNodeHeight(tree.getParent(tree.getExternalNode(index))) - tree.getNodeHeight(tree.getExternalNode(index));
            index++;
        }

        double[] vArrayOne = new double[dim];
        double[] vArrayTwo = new double[dim];

        ArrayList<Double> results = new ArrayList<Double>();

        for (Double l : lambda) {
            double distance = 0.0;
            //calculate Euclidean distance for this lambda value
            for (int i = 0; i < dim; i++) {
                vArrayOne[i] = (1.0 - l)*focalSmallM[i] + l*focalLargeM[i];
                vArrayTwo[i] = (1.0 - l)*smallMTwo[i] + l*largeMTwo[i];
                distance += Math.pow(vArrayOne[i] - vArrayTwo[i],2);
            }
            distance = Math.sqrt(distance);
            results.add(distance);
        }

        return results;

    }

    public ArrayList<Double> getMetric(Tree tree, ArrayList<Double> lambda) {

        checkTreeTaxa(focalTree, tree);

        double[] smallMTwo = new double[dim];
        double[] largeMTwo = new double[dim];

        processTree(tree, smallMTwo, largeMTwo, dim);

        double[] vArrayOne = new double[dim];
        double[] vArrayTwo = new double[dim];

        ArrayList<Double> results = new ArrayList<Double>();

        for (Double l : lambda) {
            double distance = 0.0;
            //calculate Euclidean distance for this lambda value
            for (int i = 0; i < dim; i++) {
                vArrayOne[i] = (1.0 - l)*focalSmallM[i] + l*focalLargeM[i];
                vArrayTwo[i] = (1.0 - l)*smallMTwo[i] + l*largeMTwo[i];
                distance += Math.pow(vArrayOne[i] - vArrayTwo[i],2);
            }
            distance = Math.sqrt(distance);
            results.add(distance);
        }

        return results;

    }

    /**
     * This method bypasses the constructor entirely, computing the metric on the two provided trees
     * and ignoring the internally stored tree.
     * @param tree1 Focal tree that will be used for computing the metric
     * @param tree2 Provided tree that will be compared to the focal tree
     * @param lambda Collection of lambda values for which to compute the metric
     * @return
     */
    public ArrayList<Double> getMetric_old(Tree tree1, Tree tree2, ArrayList<Double> lambda) {

        int dim = (tree1.getExternalNodeCount()-2)*(tree1.getExternalNodeCount()-1)+tree1.getExternalNodeCount();

        double[] smallMOne = new double[dim];
        double[] largeMOne = new double[dim];
        double[] smallMTwo = new double[dim];
        double[] largeMTwo = new double[dim];

        //check if taxon lists are in the same order!!
        if (tree1.getExternalNodeCount() != tree2.getExternalNodeCount()) {
            throw new RuntimeException("Different number of taxa in both trees.");
        } else {
            for (int i = 0; i < tree1.getExternalNodeCount(); i++) {
                if (!tree1.getNodeTaxon(tree1.getExternalNode(i)).getId().equals(tree2.getNodeTaxon(tree2.getExternalNode(i)).getId())) {
                    throw new RuntimeException("Mismatch between taxa in both trees: " + tree1.getNodeTaxon(tree1.getExternalNode(i)).getId() + " vs. " + tree2.getNodeTaxon(tree2.getExternalNode(i)).getId());
                }
            }
        }

        int index = 0;
        for (int i = 0; i < tree1.getExternalNodeCount(); i++) {
            for (int j = i+1; j < tree1.getExternalNodeCount(); j++) {
                //get two leaf nodes
                NodeRef nodeOne = tree1.getExternalNode(i);
                NodeRef nodeTwo = tree1.getExternalNode(j);

                //get common ancestor of 2 leaf nodes
                NodeRef MRCA = TreeUtils.getCommonAncestor(tree1, nodeOne, nodeTwo);

                int edges = 0;
                double branchLengths = 0.0;
                while (MRCA != tree1.getRoot()) {
                    edges++;
                    branchLengths += tree1.getNodeHeight(tree1.getParent(MRCA)) - tree1.getNodeHeight(MRCA);
                    MRCA = tree1.getParent(MRCA);
                }
                smallMOne[index] = edges;
                largeMOne[index] = branchLengths;
                index++;
            }
        }

        //fill out arrays further
        index = 0;
        for (int i = (tree1.getExternalNodeCount()-1)*(tree1.getExternalNodeCount()-2); i < dim; i++) {
            smallMOne[i] = 1.0;
            largeMOne[i] = tree1.getNodeHeight(tree1.getParent(tree1.getExternalNode(index))) - tree1.getNodeHeight(tree1.getExternalNode(index));
            index++;
        }

        /*for (int i = 0; i < smallMOne.length; i++) {
            System.out.print(smallMOne[i] + " ");
        }
        System.out.println();
        for (int i = 0; i < largeMOne.length; i++) {
            System.out.print(largeMOne[i] + " ");
        }
        System.out.println("\n");*/

        index = 0;
        for (int i = 0; i < tree2.getExternalNodeCount(); i++) {
            for (int j = i+1; j < tree2.getExternalNodeCount(); j++) {
                //get two leaf nodes
                NodeRef nodeOne = tree2.getExternalNode(i);
                NodeRef nodeTwo = tree2.getExternalNode(j);

                //get common ancestor of 2 leaf nodes
                NodeRef MRCA = TreeUtils.getCommonAncestor(tree2, nodeOne, nodeTwo);

                int edges = 0;
                double branchLengths = 0.0;
                while (MRCA != tree2.getRoot()) {
                    edges++;
                    branchLengths += tree2.getNodeHeight(tree2.getParent(MRCA)) - tree2.getNodeHeight(MRCA);
                    MRCA = tree2.getParent(MRCA);
                }
                smallMTwo[index] = edges;
                largeMTwo[index] = branchLengths;
                index++;
            }
        }

        //fill out arrays further
        index = 0;
        for (int i = (tree2.getExternalNodeCount()-1)*(tree2.getExternalNodeCount()-2); i < dim; i++) {
            smallMTwo[i] = 1.0;
            largeMTwo[i] = tree2.getNodeHeight(tree2.getParent(tree2.getExternalNode(index))) - tree2.getNodeHeight(tree2.getExternalNode(index));
            index++;
        }

        /*for (int i = 0; i < smallMTwo.length; i++) {
            System.out.print(smallMTwo[i] + " ");
        }
        System.out.println();
        for (int i = 0; i < largeMTwo.length; i++) {
            System.out.print(largeMTwo[i] + " ");
        }
        System.out.println("\n");*/

        double[] vArrayOne = new double[dim];
        double[] vArrayTwo = new double[dim];

        ArrayList<Double> results = new ArrayList<Double>();

        for (Double l : lambda) {
            double distance = 0.0;
            //calculate Euclidean distance for this lambda value
            for (int i = 0; i < dim; i++) {
                vArrayOne[i] = (1.0 - l)*smallMOne[i] + l*largeMOne[i];
                vArrayTwo[i] = (1.0 - l)*smallMTwo[i] + l*largeMTwo[i];
                distance += Math.pow(vArrayOne[i] - vArrayTwo[i],2);
            }
            distance = Math.sqrt(distance);
            results.add(distance);
        }

        return results;
    }

    /**
     * This method bypasses the constructor entirely, computing the metric on the two provided trees
     * and ignoring the internally stored tree.
     * @param tree1 Focal tree that will be used for computing the metric
     * @param tree2 Provided tree that will be compared to the focal tree
     * @param lambda Collection of lambda values for which to compute the metric
     * @return
     */
    public ArrayList<Double> getMetric(Tree tree1, Tree tree2, ArrayList<Double> lambda) {

        checkTreeTaxa(tree1, tree2);

        int dim = (tree1.getExternalNodeCount() - 2) * (tree1.getExternalNodeCount() - 1) + tree1.getExternalNodeCount();

        double[] smallMOne = new double[dim];
        double[] largeMOne = new double[dim];
        double[] smallMTwo = new double[dim];
        double[] largeMTwo = new double[dim];

        processTree(tree1, smallMOne, largeMOne, dim);
        processTree(tree2, smallMTwo, largeMTwo, dim);

        double[] vArrayOne = new double[dim];
        double[] vArrayTwo = new double[dim];

        ArrayList<Double> results = new ArrayList<Double>();

        for (Double l : lambda) {
            double distance = 0.0;
            //calculate Euclidean distance for this lambda value
            for (int i = 0; i < dim; i++) {
                vArrayOne[i] = (1.0 - l)*smallMOne[i] + l*largeMOne[i];
                vArrayTwo[i] = (1.0 - l)*smallMTwo[i] + l*largeMTwo[i];
                distance += Math.pow(vArrayOne[i] - vArrayTwo[i],2);
            }
            distance = Math.sqrt(distance);
            results.add(distance);
        }

        return results;
    }

    private void processTree(Tree tree, double[] smallM, double[] largeM, int dim) {

        Map<Pair<NodeRef, NodeRef>, NodeRef> permutationMap = new HashMap<Pair<NodeRef, NodeRef>, NodeRef>();
        traverse(tree, tree.getRoot(), permutationMap);

        int index = 0;

        for (int i = 0; i < tree.getExternalNodeCount(); i++) {
            for (int j = i+1; j < tree.getExternalNodeCount(); j++) {
                NodeRef MRCA = permutationMap.get(new Pair<NodeRef, NodeRef>(tree.getExternalNode(i), tree.getExternalNode(j)));

                int edges = 0;
                double branchLengths = 0.0;
                while (MRCA != tree.getRoot()) {
                    edges++;
                    branchLengths += tree.getNodeHeight(tree.getParent(MRCA)) - tree.getNodeHeight(MRCA);
                    MRCA = tree.getParent(MRCA);
                }
                smallM[index] = edges;
                largeM[index] = branchLengths;
                index++;
            }
        }

        //fill out arrays further
        index = 0;
        for (int i = (tree.getExternalNodeCount()-1)*(tree.getExternalNodeCount()-2); i < dim; i++) {
            smallM[i] = 1.0;
            largeM[i] = tree.getNodeHeight(tree.getParent(tree.getExternalNode(index))) - tree.getNodeHeight(tree.getExternalNode(index));
            index++;
        }
    }

    private void checkTreeTaxa(Tree tree1, Tree tree2) {
        //check if taxon lists are in the same order!!
        if (tree1.getExternalNodeCount() != tree2.getExternalNodeCount()) {
            throw new RuntimeException("Different number of taxa in both trees.");
        } else {
            for (int i = 0; i < tree1.getExternalNodeCount(); i++) {
                if (!tree1.getNodeTaxon(tree1.getExternalNode(i)).getId().equals(tree2.getNodeTaxon(tree2.getExternalNode(i)).getId())) {
                    throw new RuntimeException("Mismatch between taxa in both trees: " + tree1.getNodeTaxon(tree1.getExternalNode(i)).getId() + " vs. " + tree2.getNodeTaxon(tree2.getExternalNode(i)).getId());
                }
            }
        }
    }

    private Set<NodeRef> traverse(Tree tree, NodeRef node, Map<Pair<NodeRef, NodeRef>, NodeRef> permutationMap) {
        NodeRef left = tree.getChild(node, 0);
        NodeRef right = tree.getChild(node, 1);


        Set<NodeRef> leftSet = null;
        Set<NodeRef> rightSet = null;

        if (!tree.isExternal(left)) {
            leftSet = traverse(tree, left, permutationMap);
        } else {
            leftSet = Collections.singleton(left);
        }
        if (!tree.isExternal(right)) {
            rightSet = traverse(tree, right, permutationMap);
        } else {
            rightSet = Collections.singleton(right);
        }

        for (NodeRef tip1 : leftSet) {
            for (NodeRef tip2 : rightSet) {
                permutationMap.put(new Pair<NodeRef, NodeRef>(tip1, tip2), node);
            }
        }

        Set<NodeRef> tips = new HashSet<NodeRef>();
        tips.addAll(leftSet);
        tips.addAll(rightSet);

        return tips;
    }


    public static void main(String[] args) {
//        tree 1: ((A:1.2,B:0.8):0.5,(C:0.8,D:1.0):1.1);
//        tree 2: (((A:0.8,B:1.3999999999999997):0.30000000000000004,C:0.7000000000000002):0.8999999999999999,D:1.0);
//
//        lambda (0.0) = 2.0
//        lambda (0.5) = 1.9397164741270823<commit_msg>New, new algorithm. Such faster. Amaze.<commit_after>/*
 * TopologyTracer.java
 *
 * Copyright (c) 2002-2017 <NAME>, <NAME> and <NAME>
 *
 * 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.evolution.tree;

import dr.evolution.io.Importer;
import dr.evolution.io.NewickImporter;
import dr.util.Pair;

import java.io.IOException;
import java.util.*;

/**
 * @author <NAME>
 * Path difference metric according to Kendall & Colijn (2015)
 */
public class KCPathDifferenceMetric {

    private Tree focalTree;
    private int dim, externalNodeCount;
    private double[] focalSmallM, focalLargeM;

    public KCPathDifferenceMetric() {

    }

    public KCPathDifferenceMetric(Tree focalTree) {
        this.focalTree = focalTree;
        this.externalNodeCount = focalTree.getExternalNodeCount();
        this.dim = (externalNodeCount-2)*(externalNodeCount-1)+externalNodeCount;
        this.focalSmallM = new double[dim];
        this.focalLargeM = new double[dim];

        processTree(focalTree, focalSmallM, focalLargeM, dim);
    }

    public ArrayList<Double> getMetric_old(Tree tree, ArrayList<Double> lambda) {

        //check if taxon lists are in the same order!!
        if (focalTree.getExternalNodeCount() != tree.getExternalNodeCount()) {
            throw new RuntimeException("Different number of taxa in both trees.");
        } else {
            for (int i = 0; i < focalTree.getExternalNodeCount(); i++) {
                if (!focalTree.getNodeTaxon(focalTree.getExternalNode(i)).getId().equals(tree.getNodeTaxon(tree.getExternalNode(i)).getId())) {
                    throw new RuntimeException("Mismatch between taxa in both trees: " + focalTree.getNodeTaxon(focalTree.getExternalNode(i)).getId() + " vs. " + tree.getNodeTaxon(tree.getExternalNode(i)).getId());
                }
            }
        }

        double[] smallMTwo = new double[dim];
        double[] largeMTwo = new double[dim];

        int index = 0;
        for (int i = 0; i < tree.getExternalNodeCount(); i++) {
            for (int j = i+1; j < tree.getExternalNodeCount(); j++) {
                //get two leaf nodes
                NodeRef nodeOne = tree.getExternalNode(i);
                NodeRef nodeTwo = tree.getExternalNode(j);

                //get common ancestor of 2 leaf nodes
                NodeRef MRCA = TreeUtils.getCommonAncestor(tree, nodeOne, nodeTwo);

                int edges = 0;
                double branchLengths = 0.0;
                while (MRCA != tree.getRoot()) {
                    edges++;
                    branchLengths += tree.getNodeHeight(tree.getParent(MRCA)) - tree.getNodeHeight(MRCA);
                    MRCA = tree.getParent(MRCA);
                }
                smallMTwo[index] = edges;
                largeMTwo[index] = branchLengths;
                index++;
            }
        }

        //fill out arrays further
        index = 0;
        for (int i = (tree.getExternalNodeCount()-1)*(tree.getExternalNodeCount()-2); i < dim; i++) {
            smallMTwo[i] = 1.0;
            largeMTwo[i] = tree.getNodeHeight(tree.getParent(tree.getExternalNode(index))) - tree.getNodeHeight(tree.getExternalNode(index));
            index++;
        }

        double[] vArrayOne = new double[dim];
        double[] vArrayTwo = new double[dim];

        ArrayList<Double> results = new ArrayList<Double>();

        for (Double l : lambda) {
            double distance = 0.0;
            //calculate Euclidean distance for this lambda value
            for (int i = 0; i < dim; i++) {
                vArrayOne[i] = (1.0 - l)*focalSmallM[i] + l*focalLargeM[i];
                vArrayTwo[i] = (1.0 - l)*smallMTwo[i] + l*largeMTwo[i];
                distance += Math.pow(vArrayOne[i] - vArrayTwo[i],2);
            }
            distance = Math.sqrt(distance);
            results.add(distance);
        }

        return results;

    }

    public ArrayList<Double> getMetric(Tree tree, ArrayList<Double> lambda) {

        checkTreeTaxa(focalTree, tree);

        double[] smallMTwo = new double[dim];
        double[] largeMTwo = new double[dim];

        processTree(tree, smallMTwo, largeMTwo, dim);

        double[] vArrayOne = new double[dim];
        double[] vArrayTwo = new double[dim];

        ArrayList<Double> results = new ArrayList<Double>();

        for (Double l : lambda) {
            double distance = 0.0;
            //calculate Euclidean distance for this lambda value
            for (int i = 0; i < dim; i++) {
                vArrayOne[i] = (1.0 - l)*focalSmallM[i] + l*focalLargeM[i];
                vArrayTwo[i] = (1.0 - l)*smallMTwo[i] + l*largeMTwo[i];
                distance += Math.pow(vArrayOne[i] - vArrayTwo[i],2);
            }
            distance = Math.sqrt(distance);
            results.add(distance);
        }

        return results;

    }

    /**
     * This method bypasses the constructor entirely, computing the metric on the two provided trees
     * and ignoring the internally stored tree.
     * @param tree1 Focal tree that will be used for computing the metric
     * @param tree2 Provided tree that will be compared to the focal tree
     * @param lambda Collection of lambda values for which to compute the metric
     * @return
     */
    public ArrayList<Double> getMetric_old(Tree tree1, Tree tree2, ArrayList<Double> lambda) {

        int dim = (tree1.getExternalNodeCount()-2)*(tree1.getExternalNodeCount()-1)+tree1.getExternalNodeCount();

        double[] smallMOne = new double[dim];
        double[] largeMOne = new double[dim];
        double[] smallMTwo = new double[dim];
        double[] largeMTwo = new double[dim];

        //check if taxon lists are in the same order!!
        if (tree1.getExternalNodeCount() != tree2.getExternalNodeCount()) {
            throw new RuntimeException("Different number of taxa in both trees.");
        } else {
            for (int i = 0; i < tree1.getExternalNodeCount(); i++) {
                if (!tree1.getNodeTaxon(tree1.getExternalNode(i)).getId().equals(tree2.getNodeTaxon(tree2.getExternalNode(i)).getId())) {
                    throw new RuntimeException("Mismatch between taxa in both trees: " + tree1.getNodeTaxon(tree1.getExternalNode(i)).getId() + " vs. " + tree2.getNodeTaxon(tree2.getExternalNode(i)).getId());
                }
            }
        }

        int index = 0;
        for (int i = 0; i < tree1.getExternalNodeCount(); i++) {
            for (int j = i+1; j < tree1.getExternalNodeCount(); j++) {
                //get two leaf nodes
                NodeRef nodeOne = tree1.getExternalNode(i);
                NodeRef nodeTwo = tree1.getExternalNode(j);

                //get common ancestor of 2 leaf nodes
                NodeRef MRCA = TreeUtils.getCommonAncestor(tree1, nodeOne, nodeTwo);

                int edges = 0;
                double branchLengths = 0.0;
                while (MRCA != tree1.getRoot()) {
                    edges++;
                    branchLengths += tree1.getNodeHeight(tree1.getParent(MRCA)) - tree1.getNodeHeight(MRCA);
                    MRCA = tree1.getParent(MRCA);
                }
                smallMOne[index] = edges;
                largeMOne[index] = branchLengths;
                index++;
            }
        }

        //fill out arrays further
        index = 0;
        for (int i = (tree1.getExternalNodeCount()-1)*(tree1.getExternalNodeCount()-2); i < dim; i++) {
            smallMOne[i] = 1.0;
            largeMOne[i] = tree1.getNodeHeight(tree1.getParent(tree1.getExternalNode(index))) - tree1.getNodeHeight(tree1.getExternalNode(index));
            index++;
        }

        /*for (int i = 0; i < smallMOne.length; i++) {
            System.out.print(smallMOne[i] + " ");
        }
        System.out.println();
        for (int i = 0; i < largeMOne.length; i++) {
            System.out.print(largeMOne[i] + " ");
        }
        System.out.println("\n");*/

        index = 0;
        for (int i = 0; i < tree2.getExternalNodeCount(); i++) {
            for (int j = i+1; j < tree2.getExternalNodeCount(); j++) {
                //get two leaf nodes
                NodeRef nodeOne = tree2.getExternalNode(i);
                NodeRef nodeTwo = tree2.getExternalNode(j);

                //get common ancestor of 2 leaf nodes
                NodeRef MRCA = TreeUtils.getCommonAncestor(tree2, nodeOne, nodeTwo);

                int edges = 0;
                double branchLengths = 0.0;
                while (MRCA != tree2.getRoot()) {
                    edges++;
                    branchLengths += tree2.getNodeHeight(tree2.getParent(MRCA)) - tree2.getNodeHeight(MRCA);
                    MRCA = tree2.getParent(MRCA);
                }
                smallMTwo[index] = edges;
                largeMTwo[index] = branchLengths;
                index++;
            }
        }

        //fill out arrays further
        index = 0;
        for (int i = (tree2.getExternalNodeCount()-1)*(tree2.getExternalNodeCount()-2); i < dim; i++) {
            smallMTwo[i] = 1.0;
            largeMTwo[i] = tree2.getNodeHeight(tree2.getParent(tree2.getExternalNode(index))) - tree2.getNodeHeight(tree2.getExternalNode(index));
            index++;
        }

        /*for (int i = 0; i < smallMTwo.length; i++) {
            System.out.print(smallMTwo[i] + " ");
        }
        System.out.println();
        for (int i = 0; i < largeMTwo.length; i++) {
            System.out.print(largeMTwo[i] + " ");
        }
        System.out.println("\n");*/

        double[] vArrayOne = new double[dim];
        double[] vArrayTwo = new double[dim];

        ArrayList<Double> results = new ArrayList<Double>();

        for (Double l : lambda) {
            double distance = 0.0;
            //calculate Euclidean distance for this lambda value
            for (int i = 0; i < dim; i++) {
                vArrayOne[i] = (1.0 - l)*smallMOne[i] + l*largeMOne[i];
                vArrayTwo[i] = (1.0 - l)*smallMTwo[i] + l*largeMTwo[i];
                distance += Math.pow(vArrayOne[i] - vArrayTwo[i],2);
            }
            distance = Math.sqrt(distance);
            results.add(distance);
        }

        return results;
    }

    /**
     * This method bypasses the constructor entirely, computing the metric on the two provided trees
     * and ignoring the internally stored tree.
     * @param tree1 Focal tree that will be used for computing the metric
     * @param tree2 Provided tree that will be compared to the focal tree
     * @param lambda Collection of lambda values for which to compute the metric
     * @return
     */
    public ArrayList<Double> getMetric(Tree tree1, Tree tree2, ArrayList<Double> lambda) {

        checkTreeTaxa(tree1, tree2);

        int dim = (tree1.getExternalNodeCount() - 2) * (tree1.getExternalNodeCount() - 1) + tree1.getExternalNodeCount();

        double[] smallMOne = new double[dim];
        double[] largeMOne = new double[dim];
        double[] smallMTwo = new double[dim];
        double[] largeMTwo = new double[dim];

        processTree(tree1, smallMOne, largeMOne, dim);
        processTree(tree2, smallMTwo, largeMTwo, dim);

        double[] vArrayOne = new double[dim];
        double[] vArrayTwo = new double[dim];

        ArrayList<Double> results = new ArrayList<Double>();

        for (Double l : lambda) {
            double distance = 0.0;
            //calculate Euclidean distance for this lambda value
            for (int i = 0; i < dim; i++) {
                vArrayOne[i] = (1.0 - l)*smallMOne[i] + l*largeMOne[i];
                vArrayTwo[i] = (1.0 - l)*smallMTwo[i] + l*largeMTwo[i];
                distance += Math.pow(vArrayOne[i] - vArrayTwo[i],2);
            }
            distance = Math.sqrt(distance);
            results.add(distance);
        }

        return results;
    }

    private void processTree(Tree tree, double[] smallM, double[] largeM, int dim) {

        double[][] lengths = new double[tree.getExternalNodeCount()][tree.getExternalNodeCount()];
        int[][] edges = new int[tree.getExternalNodeCount()][tree.getExternalNodeCount()];

        traverse(tree, tree.getRoot(), 0.0, 0, lengths, edges);

        int index = 0;

        for (int i = 0; i < tree.getExternalNodeCount(); i++) {
            for (int j = i; j < tree.getExternalNodeCount(); j++) {
                // should work directly from 2d array
                smallM[index] = edges[i][j];
                largeM[index] = lengths[i][j];
                index++;
            }
        }
    }

    private void checkTreeTaxa(Tree tree1, Tree tree2) {
        //check if taxon lists are in the same order!!
        if (tree1.getExternalNodeCount() != tree2.getExternalNodeCount()) {
            throw new RuntimeException("Different number of taxa in both trees.");
        } else {
            for (int i = 0; i < tree1.getExternalNodeCount(); i++) {
                if (!tree1.getNodeTaxon(tree1.getExternalNode(i)).getId().equals(tree2.getNodeTaxon(tree2.getExternalNode(i)).getId())) {
                    throw new RuntimeException("Mismatch between taxa in both trees: " + tree1.getNodeTaxon(tree1.getExternalNode(i)).getId() + " vs. " + tree2.getNodeTaxon(tree2.getExternalNode(i)).getId());
                }
            }
        }
    }

    private Set<NodeRef> traverse(Tree tree, NodeRef node, double lengthFromRoot, int branchesFromRoot, double[][] lengths, int[][] edges) {
        NodeRef left = tree.getChild(node, 0);
        NodeRef right = tree.getChild(node, 1);


        Set<NodeRef> leftSet = null;
        Set<NodeRef> rightSet = null;

        if (!tree.isExternal(left)) {
            leftSet = traverse(tree, left, lengthFromRoot + tree.getBranchLength(left), branchesFromRoot + 1, lengths, edges);
        } else {
            leftSet = Collections.singleton(left);
            lengths[left.getNumber()][left.getNumber()] = tree.getBranchLength(left);
            edges[left.getNumber()][left.getNumber()] = 1;
        }
        if (!tree.isExternal(right)) {
            rightSet = traverse(tree, right, lengthFromRoot + tree.getBranchLength(right), branchesFromRoot + 1, lengths, edges);
        } else {
            rightSet = Collections.singleton(right);
            lengths[right.getNumber()][right.getNumber()] = tree.getBranchLength(right);
            edges[right.getNumber()][right.getNumber()] = 1;
        }

        for (NodeRef tip1 : leftSet) {
            for (NodeRef tip2 : rightSet) {
                if (tip1.getNumber() < tip2.getNumber()) {
                    lengths[tip1.getNumber()][tip2.getNumber()] = lengthFromRoot;
                    edges[tip1.getNumber()][tip2.getNumber()] = branchesFromRoot;
                } else {
                    lengths[tip2.getNumber()][tip1.getNumber()] = lengthFromRoot;
                    edges[tip2.getNumber()][tip1.getNumber()] = branchesFromRoot;
                }
            }
        }

        Set<NodeRef> tips = new HashSet<NodeRef>();
        tips.addAll(leftSet);
        tips.addAll(rightSet);

        return tips;
    }


    public static void main(String[] args) {
//        tree 1: ((A:1.2,B:0.8):0.5,(C:0.8,D:1.0):1.1);
//        tree 2: (((A:0.8,B:1.3999999999999997):0.30000000000000004,C:0.7000000000000002):0.8999999999999999,D:1.0);
//
//        lambda (0.0) = 2.0
//        lambda (0.5) = 1.9397164741270823
{'fraction_non_alphanumeric': 0.12222187780911938, 'fraction_numerical': 0.0193112426769164, 'mean_word_length': 2.095270075793917, 'pattern_counts': {'":': 0, '<': 107, '<?xml version=': 0, '>': 61, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 8, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '669534', 'n_tokens_mistral': 9827, 'n_tokens_neox': 9253, 'n_words': 2356}