package dr.evomodel.speciation;

import dr.evolution.coalescent.DemographicFunction;
import dr.evolution.io.NewickImporter;
import dr.evolution.tree.*;
import dr.evolution.util.MutableTaxonListListener;
import dr.evolution.util.Taxon;
import dr.evomodel.coalescent.VDdemographicFunction;
import dr.evomodel.operators.TreeNodeSlide;
import dr.evomodel.tree.TreeLogger;
import dr.evomodelxml.speciation.SpeciesTreeModelParser;
import dr.inference.model.AbstractModel;
import dr.inference.model.Model;
import dr.inference.model.Parameter;
import dr.inference.model.Variable;
import dr.inference.operators.OperatorFailedException;
import dr.inference.operators.Scalable;
import dr.util.HeapSort;
import jebl.util.FixedBitSet;

import java.util.*;

/**
 * Species tree which includes demographic function per branch.
 *
 * @author Joseph Heled
 *         Date: 24/05/2008
 */
public class SpeciesTreeModel extends AbstractModel implements MutableTree, TreeTraitProvider, TreeLogger.LogUpon, Scalable {
    private final SimpleTree spTree;
    private final SpeciesBindings species;
    private final Map<NodeRef, NodeProperties> props = new HashMap<NodeRef, NodeProperties>();
    public final Parameter sppSplitPopulations;
    private int[] singleStartPoints;
    private int[] pairStartPoints;

    private final Parameter coalPointsPops;
    private final Parameter coalPointsIndicator;
    private boolean nodePropsReady;

    private final NodeRef[] children;
    private final double[] heights;

    // any change of underlying parameters / models
    private boolean anyChange;
    // Tree has been edited in this cycle
    private boolean treeChanged;

    private final String spIndexAttrName = "spi";
    private final boolean bmp;
    private final boolean nonConstRootPopulation;
    private final boolean constantPopulation;

    private class NodeProperties {
        final int speciesIndex;
        public VDdemographicFunction demogf;
        FixedBitSet spSet;

        public NodeProperties(int n) {
            speciesIndex = n;
            demogf = null;
            spSet = new FixedBitSet(species.nSpecies());
        }
    }

    public SpeciesTreeModel(SpeciesBindings species, Parameter sppSplitPopulations,
                            Parameter coalPointsPops, Parameter coalPointsIndicator, Tree startTree,
                            boolean bmp, boolean nonConstRootPopulation, boolean constantPopulation) {
        super(SpeciesTreeModelParser.SPECIES_TREE);

        this.species = species;

        this.sppSplitPopulations = sppSplitPopulations;
        this.coalPointsPops = coalPointsPops;
        this.coalPointsIndicator = coalPointsIndicator;

        this.bmp = bmp;
        this.nonConstRootPopulation = nonConstRootPopulation;
        this.constantPopulation = constantPopulation;

        addVariable(sppSplitPopulations);

        addModel(species);

        if (coalPointsPops != null) {
            assert coalPointsIndicator != null;
            assert !constantPopulation;

            addVariable(coalPointsPops);
            addVariable(coalPointsIndicator);

            final double[][] pts = species.getPopTimesSingle();

            int start = 0;
            singleStartPoints = new int[pts.length];
            for (int i = 0; i < pts.length; i++) {
                singleStartPoints[i] = start;
                start += pts[i].length;
            }

            if (!bmp) {
                final double[][] ptp = species.getPopTimesPair();
                pairStartPoints = new int[ptp.length];
                for (int i = 0; i < ptp.length; i++) {
                    pairStartPoints[i] = start;
                    start += ptp[i].length;
                }
            }
        }

        // build an initial noninformative tree
        spTree = compatibleUninformedSpeciesTree(startTree);

        // some of the code is generic but some parts assume a binary tree.
        assert Tree.Utils.isBinary(spTree);

        final int nNodes = spTree.getNodeCount();
        heights = new double[nNodes];
        children = new NodeRef[2 * nNodes + 1];

        // fixed properties
        for (int k = 0; k < getExternalNodeCount(); ++k) {
            final NodeRef nodeRef = getExternalNode(k);
            final int n = (Integer) getNodeAttribute(nodeRef, spIndexAttrName);
            final NodeProperties np = new NodeProperties(n);
            props.put(nodeRef, np);
            np.spSet.set(n);
        }

        for (int k = 0; k < getInternalNodeCount(); ++k) {
            final NodeRef nodeRef = getInternalNode(k);
            props.put(nodeRef, new NodeProperties(-1));
        }

        nodePropsReady = false;

        // crappy way to pass a result back from compatibleUninformedSpeciesTree.
        // check is using isCompatible(), which requires completion of construction.
        boolean check = spTree.getAttribute("check") != null;
        spTree.setAttribute("check", null);

        if (check) {
            // change only is needed - if the user provided a compatible state she may know
            // what she is doing
            for (SpeciesBindings.GeneTreeInfo gt : species.getGeneTrees()) {
                if (!isCompatible(gt)) {

                    species.makeCompatible(spTree.getRootHeight());
                    for (SpeciesBindings.GeneTreeInfo t : species.getGeneTrees()) {
                        assert isCompatible(t);
                    }
                    anyChange = false;
                    break;
                }
            }
        }
    }

    public boolean constPopulation() {
        return constantPopulation;
    }

    // Is gene tree compatible with species tree

    public boolean isCompatible(SpeciesBindings.GeneTreeInfo geneTreeInfo) {
        // can't set demographics if a tree is not compatible, but we need spSets.
        if (!nodePropsReady) {
            setSPsets(getRoot());
        }
        return isCompatible(getRoot(), geneTreeInfo.getCoalInfo(), 0) >= 0;
    }

    // Not very efficient, should do something better, based on traversing the cList once

    private int isCompatible(NodeRef node, SpeciesBindings.CoalInfo[] cList, int loc) {
        if (!isExternal(node)) {
            int l = -1;
            for (int nc = 0; nc < getChildCount(node); ++nc) {
                int l1 = isCompatible(getChild(node, nc), cList, loc);
                if (l1 < 0) {
                    return -1;
                }
                assert l == -1 || l1 == l;

                l = l1;
            }
            loc = l;

            assert cList[loc].ctime >= getNodeHeight(node);
        }

        if (node == getRoot()) {
            return cList.length;
        }

        // spSet guaranteed to be ready by caller
        final FixedBitSet nodeSps = props.get(node).spSet;

        final double limit = getNodeHeight(getParent(node));

        while (loc < cList.length) {
            final SpeciesBindings.CoalInfo ci = cList[loc];
            if (ci.ctime >= limit) {
                break;
            }
            boolean allIn = true, noneIn = true;

            for (int i = 0; i < 2; ++i) {
                final FixedBitSet s = ci.sinfo[i];

                final int in1 = s.intersectCardinality(nodeSps);
                if (in1 > 0) {
                    noneIn = false;
                }
                if (s.cardinality() != in1) {
                    allIn = false;
                }
            }
            if (!(allIn || noneIn)) {
                return -1;
            }
            ++loc;
        }
        return loc;
    }


    private static double
    fp(double val, double low, double[][] tt, int[] ii) {
        for (int k = 0; k < ii.length; ++k) {
            int ip = ii[k];
            if (ip == tt[k].length || val <= tt[k][ip]) {
                --ip;
                while (ip >= 0 && val <= tt[k][ip]) {
                    --ip;
                }
                assert ((ip < 0) || (tt[k][ip] < val)) && ((ip + 1 == tt[k].length) || (val <= tt[k][ip + 1]));
                if (ip >= 0) {
                    low = Math.max(low, tt[k][ip]);
                }
            } else {
                ++ip;
                while (ip < tt[k].length && val > tt[k][ip]) {
                    ++ip;
                }
                assert tt[k][ip - 1] < val && ((ip == tt[k].length) || (val <= tt[k][ip]));
                low = Math.max(low, tt[k][ip - 1]);
            }
        }
        return low;
    }

    private interface SimpleDemographicFunction {
        double population(double t);

        double upperBound();
    }

    private class PLSD implements SimpleDemographicFunction {
        private final double[] pops;
        private final double[] times;

        public PLSD(double[] pops, double[] times) {
            assert pops.length == times.length + 1;

            this.pops = pops;
            this.times = times;
        }

        public double population(double t) {
            if (t >= upperBound()) {
                return pops[pops.length - 1];
            }

            int k = 0;
            while (t > times[k]) {
                t -= times[k];
                ++k;
            }
            double a = t / (times[k] - (k > 0 ? times[k - 1] : 0));
            return a * pops[k] + (1 - a) * pops[k + 1];
        }

        public double upperBound() {
            return times[times.length - 1];
        }
    }
    // Pass arguments of recursive functions in a compact format.

    private class Args {
        final double[][] cps = species.getPopTimesSingle();
        final double[][] cpp; // = species.getPopTimesPair();
        final int[] iSingle = new int[cps.length];
        final int[] iPair; // = new int[cpp.length];

        final double[] indicators = ((Parameter.Default) coalPointsIndicator).inspectParameterValues();
        final double[] pops = ((Parameter.Default) coalPointsPops).inspectParameterValues();
        final SimpleDemographicFunction[] dms;

        Args(Boolean bmp) {
            if (!bmp) {
                cpp = species.getPopTimesPair();
                iPair = new int[cpp.length];
                dms = null;
            } else {
                cpp = null;
                iPair = null;
                int nsps = cps.length;
                dms = new SimpleDemographicFunction[nsps];
                for (int nsp = 0; nsp < nsps; ++nsp) {
                    final int start = singleStartPoints[nsp];
                    final int stop = nsp < nsps - 1 ? singleStartPoints[nsp + 1] : pops.length;

                    double[] pop = new double[1 + stop - start];
                    pop[0] = sppSplitPopulations.getParameterValue(nsp); //  pops[nsp];

                    for (int k = 0; k < stop - start; ++k) {
                        pop[k + 1] = pops[start + k];
                    }
                    dms[nsp] = new PLSD(pop, cps[nsp]);
                }
            }
        }

        private double findPrev(double val, double low) {
            low = fp(val, low, cps, iSingle);
            low = fp(val, low, cpp, iPair);

            return low;
        }
    }

    class RawPopulationHelper {
        final int[] preOrderIndices = new int[getNodeCount()];
        final double[] pops = ((Parameter.Default) sppSplitPopulations).inspectParameterValues();
        final int nsp = species.nSpecies();
        final Args args = coalPointsPops != null ? new Args(bmp) : null;

        RawPopulationHelper() {
            setPreorderIndices(preOrderIndices);
        }

        public void getPopulations(NodeRef n, int nc, double[] p) {
            p[1] = pops[nsp + 2 * preOrderIndices[n.getNumber()] + nc];

            final NodeRef child = getChild(n, nc);
            if (isExternal(child)) {
                p[0] = pops[props.get(child).speciesIndex];
            } else {
                int k = nsp + 2 * preOrderIndices[child.getNumber()];
                p[0] = pops[k] + pops[k + 1];
            }
        }

        public double tipPopulation(NodeRef tip) {
            return pops[props.get(tip).speciesIndex];
        }

        public int nSpecies() {
            return species.nSpecies();
        }

        public boolean perSpeciesPopulation() {
            return args != null;
        }

        public double[] getTimes(int ns) {
            return ((PLSD) args.dms[ns]).times;
        }

        public double[] getPops(int ns) {
            return ((PLSD) args.dms[ns]).pops;
        }

        public void getRootPopulations(double[] p) {
            int k = nsp + 2 * preOrderIndices[getRoot().getNumber()];
            p[0] = pops[k] + pops[k + 1];
            p[1] = nonConstRootPopulation ? pops[pops.length - 1] : p[0];
        }

        public double geneTreesRootHeight() {
            //getNodeDemographic(getRoot()).
            double h = -1;
            for (SpeciesBindings.GeneTreeInfo t : species.getGeneTrees()) {
                h = Math.max(h, t.tree.getNodeHeight(t.tree.getRoot()));
            }
            return h;
        }
    }

    RawPopulationHelper getPopulationHelper() {
        return new RawPopulationHelper();
    }

    static private class Points implements Comparable<Points> {
        final double time;
        double population;
        final boolean use;

        Points(double t, double p) {
            time = t;
            population = p;
            use = true;
        }

        Points(double t, boolean u) {
            time = t;
            population = 0;
            use = u;
        }

        public int compareTo(Points points) {
            return time < points.time ? -1 : (time > points.time ? 1 : 0);
        }
    }

    private NodeProperties
    setSPsets(NodeRef nodeID) {
        final NodeProperties nprop = props.get(nodeID);

        if (!isExternal(nodeID)) {
            nprop.spSet = new FixedBitSet(species.nSpecies());
            for (int nc = 0; nc < getChildCount(nodeID); ++nc) {
                NodeProperties p = setSPsets(getChild(nodeID, nc));
                nprop.spSet.union(p.spSet);
            }
        }
        return nprop;
    }

    private int ti2f(int i, int j) {
        return (i == 0) ? j : 2 * i + j + 1;
    }

    private VDdemographicFunction
    bestLinearFit(double[] xs, double[] ys, boolean[] use) {

        assert (xs.length + 1) == ys.length;
        assert ys.length == use.length + 2 || ys.length == use.length + 1;

        int N = ys.length;
        if (N == 2) {
            // cheaper
            return new VDdemographicFunction(xs, ys, getUnits());
        }

        List<Integer> iv = new ArrayList<Integer>(2);
        iv.add(0);
        for (int k = 0; k < N - 2; ++k) {
            if (use[k]) {
                iv.add(k + 1);
            }
        }
        iv.add(N - 1);

        double[] ati = new double[xs.length + 1];
        ati[0] = 0.0;
        System.arraycopy(xs, 0, ati, 1, xs.length);
        int n = iv.size();

        double[] a = new double[3 * n];
        double[] v = new double[n];

        for (int k = 0; k < n - 1; ++k) {
            int i0 = iv.get(k);
            int i1 = iv.get(k + 1);

            double u0 = ati[i0];
            double u1 = ati[i1] - ati[i0];
            // on last interval add data for last point
            if (i1 == N - 1) {
                i1 += 1;
            }

            final int l = ti2f(k, k);
            final int l1 = ti2f(k + 1, k);

            for (int j = i0; j < i1; ++j) {
                double t = ati[j];
                double y = ys[j];

                double z = (t - u0) / u1;
                v[k] += y * (1 - z);

                a[l] += (1 - z) * (1 - z);
                a[l + 1] += z * (1 - z);

                a[l1] += z * (1 - z);
                a[l1 + 1] += z * z;
                v[k + 1] += y * z;
            }
        }

        for (int k = 0; k < n - 1; ++k) {

            final double r = a[ti2f(k + 1, k)] / a[ti2f(k, k)];
            for (int j = k; j < k + 3; ++j) {
                a[ti2f((k + 1), j)] -= a[ti2f(k, j)] * r;
            }
            v[k + 1] -= v[k] * r;
        }

        double[] z = new double[n];
        for (int k = n - 1; k > 0; --k) {
            z[k] = v[k] / a[ti2f(k, k)];
            v[k - 1] -= a[ti2f((k - 1), k)] * z[k];
        }

        z[0] = v[0] / a[ti2f(0, 0)];
        double[] t = new double[iv.size() - 1];
        for (int j = 0; j < t.length; ++j) {
            t[j] = ati[iv.get(j + 1)];
        }
        return new VDdemographicFunction(t, z, getUnits());
    }

    //  Assign positions in 'pointsList' for the sub-tree rooted at the ancestor of
    //  nodeID.
    //
    //  pointsList is indexed by node-id. Every element is a list of internal
    //  population points for the branch between nodeID and it's ancestor
    //

    private NodeProperties getDemographicPoints(final NodeRef nodeID, Args args, Points[][] pointsList) {

        final NodeProperties nprop = props.get(nodeID);
        final int nSpecies = species.nSpecies();

        // Species assignment from the tips never changes
        if (!isExternal(nodeID)) {
            nprop.spSet = new FixedBitSet(nSpecies);
            for (int nc = 0; nc < getChildCount(nodeID); ++nc) {
                final NodeProperties p = getDemographicPoints(getChild(nodeID, nc), args, pointsList);
                nprop.spSet.union(p.spSet);
            }
        }

        if (args == null) {
            return nprop;
        }

        // parent height
        final double cHeight = nodeID != getRoot() ? getNodeHeight(getParent(nodeID)) : Double.MAX_VALUE;

        // points along branch
        // not sure what a good default size is?
        List<Points> allPoints = new ArrayList<Points>(5);

        if (bmp) {
            for (int isp = nprop.spSet.nextOnBit(0); isp >= 0; isp = nprop.spSet.nextOnBit(isp + 1)) {
                final double[] cp = args.cps[isp];
                final int upi = singleStartPoints[isp];

                int i = args.iSingle[isp];
                for (/**/; i < cp.length && cp[i] < cHeight; ++i) {
                    allPoints.add(new Points(cp[i], args.indicators[upi + i] > 0));
                }
                args.iSingle[isp] = i;
            }
        } else {

            for (int isp = nprop.spSet.nextOnBit(0); isp >= 0; isp = nprop.spSet.nextOnBit(isp + 1)) {
                final double nodeHeight = spTree.getNodeHeight(nodeID);
                {
                    double[] cp = args.cps[isp];
                    final int upi = singleStartPoints[isp];

                    int i = args.iSingle[isp];

                    while (i < cp.length && cp[i] < cHeight) {
                        if (args.indicators[upi + i] > 0) {
                            //System.out.println("  popbit s");
                            args.iSingle[isp] = i;
                            double prev = args.findPrev(cp[i], nodeHeight);
                            double mid = (prev + cp[i]) / 2.0;
                            assert nodeHeight < mid;
                            allPoints.add(new Points(mid, args.pops[upi + i]));
                        }
                        ++i;
                    }
                    args.iSingle[isp] = i;
                }

                final int kx = (isp * (2 * nSpecies - isp - 3)) / 2 - 1;
                for (int y = nprop.spSet.nextOnBit(isp + 1); y >= 0; y = nprop.spSet.nextOnBit(y + 1)) {

                    assert isp < y;
                    int k = kx + y;

                    double[] cp = args.cpp[k];
                    int i = args.iPair[k];
                    final int upi = pairStartPoints[k];

                    while (i < cp.length && cp[i] < cHeight) {
                        if (args.indicators[upi + i] > 0) {
                            //System.out.println("  popbit p");
                            args.iPair[k] = i;
                            final double prev = args.findPrev(cp[i], nodeHeight);
                            double mid = (prev + cp[i]) / 2.0;
                            assert nodeHeight < mid;
                            allPoints.add(new Points(mid, args.pops[upi + i]));
                        }
                        ++i;
                    }
                    args.iPair[k] = i;
                }
            }
        }

        Points[] all = null;

        if (allPoints.size() > 0) {
            all = allPoints.toArray(new Points[allPoints.size()]);
            if (all.length > 1) {
                HeapSort.sort(all);
            }

            int len = all.length;

            if (bmp) {
                int k = 0;
                while (k + 1 < len) {
                    final double t = all[k].time;
                    if (t == all[k + 1].time) {
                        int j = k + 2;
                        boolean use = all[k].use || all[k + 1].use;

                        while (j < len && t == all[j].time) {
                            use = use || all[j].use;
                            j += 1;
                        }
                        int removed = (j - k - 1);
                        all[k] = new Points(t, use);
                        for (int i = k + 1; i < len - removed; ++i) {
                            all[i] = all[i + removed];
                        }
                        len -= removed;
                    }
                    ++k;
                }
            } else {
                // duplications

                int k = 0;
                while (k + 1 < len) {
                    double t = all[k].time;
                    if (t == all[k + 1].time) {
                        int j = k + 2;
                        double v = all[k].population + all[k + 1].population;
                        while (j < len && t == all[j].time) {
                            v += all[j].population;
                            j += 1;
                        }
                        int removed = (j - k - 1);
                        all[k] = new Points(t, v / (removed + 1));
                        for (int i = k + 1; i < len - removed; ++i) {
                            all[i] = all[i + removed];
                        }
                        //System.arraycopy(all, j, all, k + 1, all.length - j + 1);
                        len -= removed;
                    }
                    ++k;
                }
            }

            if (len != all.length) {
                Points[] a = new Points[len];
                System.arraycopy(all, 0, a, 0, len);
                all = a;
            }

            if (bmp) {
                for (Points p : all) {
                    double t = p.time;
                    assert p.population == 0;
                    for (int isp = nprop.spSet.nextOnBit(0); isp >= 0; isp = nprop.spSet.nextOnBit(isp + 1)) {
                        SimpleDemographicFunction d = args.dms[isp];
                        if (t <= d.upperBound()) {
                            p.population += d.population(t);
                        }
                    }
                }
            }
        }

        pointsList[nodeID.getNumber()] = all;

        return nprop;
    }

    private int setDemographics(NodeRef nodeID, int pStart, int side, double[] pops, Points[][] pointsList) {

        final int nSpecies = species.nSpecies();
        final NodeProperties nprop = props.get(nodeID);
        int pEnd;
        double p0;

        if (isExternal(nodeID)) {
            final int sps = nprop.speciesIndex;
            p0 = pops[sps];
            pEnd = pStart;
        } else {
            assert getChildCount(nodeID) == 2;

            final int iHere = setDemographics(getChild(nodeID, 0), pStart, 0, pops, pointsList);
            pEnd = setDemographics(getChild(nodeID, 1), iHere + 1, 1, pops, pointsList);
            if (constantPopulation) {
                final int i = nSpecies + iHere;
                p0 = pops[i];
            } else {
                final int i = nSpecies + iHere * 2;
                p0 = pops[i] + pops[i + 1];
            }
        }

        if (constantPopulation) {
            double[] xs = {};
            double[] ys = {p0};

            nprop.demogf = new VDdemographicFunction(xs, ys, getUnits());
            // new ConstantPopulation(p0, getUnits());
        } else {
            final double t0 = getNodeHeight(nodeID);

            Points[] p = pointsList != null ? pointsList[nodeID.getNumber()] : null;
            final int plen = p == null ? 0 : p.length;

            final boolean isRoot = nodeID == getRoot();
//        double[] xs = new double[plen + (isRoot ? 1 : 1)];
//        double[] ys = new double[plen + (isRoot ? 2 : 2)];

            final boolean useBMP = bmp && pointsList != null;
            // internal nodes add one population point for the branch end.
            // on the root (with bmp) there is no such point.
            final int len = plen + (useBMP ? (!isRoot ? 1 : 0) : 1);
            double[] xs = new double[len];
            double[] ys = new double[len + 1];
            boolean[] use = new boolean[len];
            ys[0] = p0;
            for (int i = 0; i < plen; ++i) {
                xs[i] = p[i].time - t0;
                ys[i + 1] = p[i].population;
                use[i] = p[i].use;
            }

            if (!isRoot) {
                final int anccIndex = (side == 0) ? pEnd : pStart - 1;
                final double pe = pops[nSpecies + anccIndex * 2 + side];
                final double b = getBranchLength(nodeID);

                xs[xs.length - 1] = b;
                ys[ys.length - 1] = pe;
            }

            if (useBMP) {
                nprop.demogf = bestLinearFit(xs, ys, use);
            } else {

                if (isRoot) {
                    // extend the last point to most ancient coalescent point. Has no effect on the demographic
                    // per se but for use when analyzing the results.

                    double h = -1;
                    for (SpeciesBindings.GeneTreeInfo t : species.getGeneTrees()) {
                        h = Math.max(h, t.tree.getNodeHeight(t.tree.getRoot()));
                    }

                    final double rh = h - t0;
                    xs[xs.length - 1] = rh; //getNodeHeight(nodeID);
                    //spTree.setBranchLength(nodeID, rh);

                    // last value is for root branch end point
                    ys[ys.length - 1] = pointsList != null ? ys[ys.length - 2] :
                            (nonConstRootPopulation ? pops[pops.length - 1] : ys[ys.length - 2]);
                }

                nprop.demogf = new VDdemographicFunction(xs, ys, getUnits());
            }
        }
        return pEnd;
    }

    private void setNodeProperties() {
        Points[][] perBranchPoints = null;

        if (coalPointsPops != null) {
            final Args args = new Args(bmp);
            perBranchPoints = new Points[getNodeCount()][];
            getDemographicPoints(getRoot(), args, perBranchPoints);
        } else {
            // sets species info
            getDemographicPoints(getRoot(), null, null);
        }

        setDemographics(getRoot(), 0, -1, ((Parameter.Default) sppSplitPopulations).inspectParameterValues(), perBranchPoints);
    }

    private Map<NodeRef, NodeProperties> getProps() {
        if (!nodePropsReady) {
            setNodeProperties();
            nodePropsReady = true;
        }
        return props;
    }

    public DemographicFunction getNodeDemographic(NodeRef node) {
        return getProps().get(node).demogf;
    }

    public FixedBitSet spSet(NodeRef node) {
        return getProps().get(node).spSet;
    }

    public int speciesIndex(NodeRef tip) {
        assert isExternal(tip);

        // always ready even if props is dirty
        return props.get(tip).speciesIndex;
    }

    private Double setInitialSplitPopulations(FlexibleTree startTree, NodeRef node, int pos[]) {
        if (!startTree.isExternal(node)) {
            int loc = -1;
            for (int nc = 0; nc < startTree.getChildCount(node); ++nc) {
                final Double p = setInitialSplitPopulations(startTree, startTree.getChild(node, nc), pos);
                if (!constantPopulation && nc == 0) {
                    loc = pos[0];
                    pos[0] += 1;
                }
                if (p != null) {
                    if (constantPopulation) {
                        //
                    } else {
                        sppSplitPopulations.setParameterValueQuietly(species.nSpecies() + 2 * loc + nc, p);
                    }
                }
            }
        }

        final String comment = (String) startTree.getNodeAttribute(node, NewickImporter.COMMENT);
        Double p0 = null;
        if (comment != null) {
            StringTokenizer st = new StringTokenizer(comment);

            p0 = Double.parseDouble(st.nextToken());
            if (startTree.isExternal(node)) {
                int ns = (Integer) startTree.getNodeAttribute(node, spIndexAttrName);
                sppSplitPopulations.setParameterValueQuietly(ns, p0);
            } else if (constantPopulation) {
                // not tested code !!
                sppSplitPopulations.setParameterValueQuietly(species.nSpecies() + pos[0], p0);
                pos[0] += 1;
            }

            // if just one value const
            if (st.hasMoreTokens()) {
                p0 = Double.parseDouble(st.nextToken());
            }
        }
        return !constantPopulation ? p0 : null;
    }

    private SimpleTree compatibleUninformedSpeciesTree(Tree startTree) {
        double rootHeight = Double.MAX_VALUE;

        for (SpeciesBindings.GeneTreeInfo t : species.getGeneTrees()) {
            rootHeight = Math.min(rootHeight, t.getCoalInfo()[0].ctime);
        }

        final SpeciesBindings.SPinfo[] spp = species.species;

        if (startTree != null) {
            // Allow start tree to be very basic basic - may be only partially resolved and no
            // branch lengths

            if (startTree.getExternalNodeCount() != spp.length) {
                throw new Error("Start tree error - different number of tips");
            }

            final FlexibleTree tree = new FlexibleTree(startTree, true);
            tree.resolveTree();
            final double treeHeight = tree.getRootHeight();
            if (treeHeight <= 0) {
                tree.setRootHeight(1.0);
                Utils.correctHeightsForTips(tree);
                SimpleTree.Utils.scaleNodeHeights(tree, rootHeight / tree.getRootHeight());
            }

            SimpleTree sTree = new SimpleTree(tree);
            for (int ns = 0; ns < spp.length; ns++) {
                SpeciesBindings.SPinfo sp = spp[ns];
                final int i = sTree.getTaxonIndex(sp.name);
                if (i < 0) {
                    throw new Error(sp.name + " is not present in the start tree");
                }
                final SimpleNode node = sTree.getExternalNode(i);
                node.setAttribute(spIndexAttrName, ns);

                // set for possible pops
                tree.setNodeAttribute(tree.getNode(tree.getTaxonIndex(sp.name)), spIndexAttrName, ns);
            }

            if (treeHeight > 0) {
                sTree.setAttribute("check", new Double(rootHeight));
            }

            {
                //assert ! constantPopulation; // not implemented yet
                int[] pos = {0};
                setInitialSplitPopulations(tree, tree.getRoot(), pos);
            }

            return sTree;
        }


        final double delta = rootHeight / (spp.length + 1);
        double cTime = delta;

        List<SimpleNode> subs = new ArrayList<SimpleNode>(spp.length);

        for (int ns = 0; ns < spp.length; ns++) {
            SpeciesBindings.SPinfo sp = spp[ns];
            final SimpleNode node = new SimpleNode();
            node.setTaxon(new Taxon(sp.name));
            subs.add(node);

            node.setAttribute(spIndexAttrName, ns);
        }

        while (subs.size() > 1) {
            final SimpleNode node = new SimpleNode();
            int i = 0, j = 1;
            node.addChild(subs.get(i));
            node.addChild(subs.get(j));
            node.setHeight(cTime);
            cTime += delta;
            subs.set(j, node);
            subs.remove(i);
        }

        return new SimpleTree(subs.get(0));
    }

    public void setPreorderIndices(int[] indices) {
        setPreorderIndices(getRoot(), 0, indices);
    }

    private int setPreorderIndices(NodeRef node, int loc, int[] indices) {
        if (!isExternal(node)) {
            int l = setPreorderIndices(getChild(node, 0), loc, indices);
            indices[node.getNumber()] = l;
            loc = setPreorderIndices(getChild(node, 1), l + 1, indices);
        }
        return loc;
    }

    public String getName() {
        return getModelName();
    }

    static private TreeNodeSlide internalTreeOP = null;

    public int scale(double scaleFactor, int nDims) throws OperatorFailedException {
        assert scaleFactor > 0;
        if (nDims <= 0) {
            // actually when in an up down with operators on the gene trees the flags
            // may indicate a change
            //storeState();  // just checks assert really
            beginTreeEdit();
            final int count = getInternalNodeCount();
            for (int i = 0; i < count; ++i) {
                final NodeRef n = getInternalNode(i);
                setNodeHeight(n, getNodeHeight(n) * scaleFactor);
            }
            endTreeEdit();
            fireModelChanged(this, 1);
            return count;
        } else {
            if (nDims != 1) {
                throw new OperatorFailedException("not implemented for count != 1");
            }
            if (internalTreeOP == null) {
                internalTreeOP = new TreeNodeSlide(this, species, 1);
            }

            internalTreeOP.operateOneNode(scaleFactor);
            fireModelChanged(this, 1);
            return nDims;
        }
    }

    private final boolean verbose = false;

    protected void handleModelChangedEvent(Model model, Object object, int index) {
        if (verbose) System.out.println(" SPtree: model changed " + model.getId());

        nodePropsReady = false;
        anyChange = true;
        // this should happen by default, no?
        fireModelChanged();
    }

    protected final void handleVariableChangedEvent(Variable variable, int index, Parameter.ChangeType type) {
        if (verbose) System.out.println(" SPtree: parameter changed " + variable.getId());

        nodePropsReady = false;
        anyChange = true;
    }

    protected void storeState() {
        assert !treeChanged;
        assert !anyChange;
    }

    protected void restoreState() {
        if (verbose) System.out.println(" SPtree: restore (" + treeChanged + "," + anyChange + ")");

        if (treeChanged) {
            //
            spTree.beginTreeEdit();

            for (int k = 0; k < getInternalNodeCount(); ++k) {
                final NodeRef node = getInternalNode(k);
                final int index = node.getNumber();
                final double h = heights[index];
                if (getNodeHeight(node) != h) {
                    setNodeHeight(node, h);
                }
                for (int nc = 0; nc < 2; ++nc) {
                    final NodeRef child = getChild(node, nc);

                    final NodeRef child1 = children[2 * index + nc];
                    if (child != child1) {
                        replaceChild(node, child, child1);
                    }
                    assert getParent(child1) == node;
                }
            }
            setRoot(children[children.length - 1]);

            if (verbose) System.out.println("  restored to: " + spTree);

            spTree.endTreeEdit();
        }
        if (treeChanged || anyChange) {
            setNodeProperties();
        }
        treeChanged = false;
        anyChange = false;
    }

    protected void acceptState() {
        if (verbose) System.out.println(" SPtree: accept");

        treeChanged = false;
        anyChange = false;
    }

    String previousTopology = null;

    public boolean logNow(long state) {
        final String curTop = Tree.Utils.uniqueNewick(spTree, spTree.getRoot());
        if (state == 0 || !curTop.equals(previousTopology)) {
            previousTopology = curTop;
            return true;
        }
        return false;
    }

//    TreeTrait dmf = new TreeTrait.S() {
//        public String getTraitName() {
//            return "dmf";
//        }
//
//        public Intent getIntent() {
//            return Intent.NODE;
//        }
//
//        public String getTrait(Tree tree, NodeRef node) {
//            assert tree == SpeciesTreeModel.this;
//
//            //final VDdemographicFunction df = getProps().get(node).demogf;
//
//            final DemographicFunction df = getNodeDemographic(node);
//            return df.toString();
//        }
//    };

    TreeTrait dmt = new TreeTrait.DA() {
        public String getTraitName() {
            return "dmt";
        }

        public Intent getIntent() {
            return Intent.NODE;
        }

        public double[] getTrait(Tree tree, NodeRef node) {
            assert tree == SpeciesTreeModel.this;

            final VDdemographicFunction df = (VDdemographicFunction) getNodeDemographic(node);
            return df.times();
        }
    };

    TreeTrait dmv = new TreeTrait.DA() {
        public String getTraitName() {
            return "dmv";
        }

        public Intent getIntent() {
            return Intent.NODE;
        }

        public double[] getTrait(Tree tree, NodeRef node) {
            assert tree == SpeciesTreeModel.this;

            final VDdemographicFunction df = (VDdemographicFunction) getNodeDemographic(node);
            return df.values();
        }
    };


    public TreeTrait[] getTreeTraits() {
        return new TreeTrait[]{dmt, dmv};
    }

    public TreeTrait getTreeTrait(String key) {
        if (key.equals(dmt.getTraitName())) {
            return dmt;
        } else if (key.equals(dmv.getTraitName())) {
            return dmv;
        }
        throw new IllegalArgumentException();
    }


    // boring delegation

    public SimpleTree getSimpleTree() {
        return spTree;
    }

    public Tree getCopy() {
        return spTree.getCopy();
    }

    public Type getUnits() {
        return spTree.getUnits();
    }

    public void setUnits(Type units) {
        spTree.setUnits(units);
    }

    public int getNodeCount() {
        return spTree.getNodeCount();
    }

    public boolean hasNodeHeights() {
        return spTree.hasNodeHeights();
    }

    public double getNodeHeight(NodeRef node) {
        return spTree.getNodeHeight(node);
    }

    public double getNodeRate(NodeRef node) {
        return spTree.getNodeRate(node);
    }

    public Taxon getNodeTaxon(NodeRef node) {
        return spTree.getNodeTaxon(node);
    }

    public int getChildCount(NodeRef node) {
        return spTree.getChildCount(node);
    }

    public boolean isExternal(NodeRef node) {
        return spTree.isExternal(node);
    }

    public boolean isRoot(NodeRef node) {
        return spTree.isRoot(node);
    }

    public NodeRef getChild(NodeRef node, int i) {
        return spTree.getChild(node, i);
    }

    public NodeRef getParent(NodeRef node) {
        return spTree.getParent(node);
    }

    public boolean hasBranchLengths() {
        return spTree.hasBranchLengths();
    }

    public double getBranchLength(NodeRef node) {
        return spTree.getBranchLength(node);
    }

    public void setBranchLength(NodeRef node, double length) {
        spTree.setBranchLength(node, length);
    }

    public NodeRef getExternalNode(int i) {
        return spTree.getExternalNode(i);
    }

    public NodeRef getInternalNode(int i) {
        return spTree.getInternalNode(i);
    }

    public NodeRef getNode(int i) {
        return spTree.getNode(i);
    }

    public int getExternalNodeCount() {
        return spTree.getExternalNodeCount();
    }

    public int getInternalNodeCount() {
        return spTree.getInternalNodeCount();
    }

    public NodeRef getRoot() {
        return spTree.getRoot();
    }

    public void setRoot(NodeRef r) {
        spTree.setRoot(r);
    }

    public void addChild(NodeRef p, NodeRef c) {
        spTree.addChild(p, c);
    }

    public void removeChild(NodeRef p, NodeRef c) {
        spTree.removeChild(p, c);
    }

    public void replaceChild(NodeRef node, NodeRef child, NodeRef newChild) {
        spTree.replaceChild(node, child, newChild);
    }

    public boolean beginTreeEdit() {
        boolean beingEdited = spTree.beginTreeEdit();
        if (!beingEdited) {
            // save tree for restore
            for (int n = 0; n < getInternalNodeCount(); ++n) {
                final NodeRef node = getInternalNode(n);
                final int k = node.getNumber();
                children[2 * k] = getChild(node, 0);
                children[2 * k + 1] = getChild(node, 1);
                heights[k] = getNodeHeight(node);
            }
            children[children.length - 1] = getRoot();

            treeChanged = true;
            nodePropsReady = false;
            //anyChange = true;
        }
        return beingEdited;
    }

    public void endTreeEdit() {
        spTree.endTreeEdit();
        fireModelChanged();
    }

    public void setNodeHeight(NodeRef n, double height) {
        spTree.setNodeHeight(n, height);
    }

    public void setNodeRate(NodeRef n, double rate) {
        spTree.setNodeRate(n, rate);
    }

    public void setNodeAttribute(NodeRef node, String name, Object value) {
        spTree.setNodeAttribute(node, name, value);
    }

    public Object getNodeAttribute(NodeRef node, String name) {
        return spTree.getNodeAttribute(node, name);
    }

    public Iterator getNodeAttributeNames(NodeRef node) {
        return spTree.getNodeAttributeNames(node);
    }

    public int getTaxonCount() {
        return spTree.getTaxonCount();
    }

    public Taxon getTaxon(int taxonIndex) {
        return spTree.getTaxon(taxonIndex);
    }

    public String getTaxonId(int taxonIndex) {
        return spTree.getTaxonId(taxonIndex);
    }

    public int getTaxonIndex(String id) {
        return spTree.getTaxonIndex(id);
    }

    public int getTaxonIndex(Taxon taxon) {
        // can't compare taxa
        return getTaxonIndex(taxon.getId());
    }

    public List<Taxon> asList() {
        return spTree.asList();
    }

    public Iterator<Taxon> iterator() {
        return spTree.iterator();
    }

    public Object getTaxonAttribute(int taxonIndex, String name) {
        return spTree.getTaxonAttribute(taxonIndex, name);
    }

    public int addTaxon(Taxon taxon) {
        return spTree.addTaxon(taxon);
    }

    public boolean removeTaxon(Taxon taxon) {
        return spTree.removeTaxon(taxon);
    }

    public void setTaxonId(int taxonIndex, String id) {
        spTree.setTaxonId(taxonIndex, id);
    }

    public void setTaxonAttribute(int taxonIndex, String name, Object value) {
        spTree.setTaxonAttribute(taxonIndex, name, value);
    }

    public String getId() {
        return spTree.getId();
    }

    public void setId(String id) {
        spTree.setId(id);
    }

    public void setAttribute(String name, Object value) {
        spTree.setAttribute(name, value);
    }

    public Object getAttribute(String name) {
        return spTree.getAttribute(name);
    }

    public Iterator<String> getAttributeNames() {
        return spTree.getAttributeNames();
    }

    public void addMutableTreeListener(MutableTreeListener listener) {
        spTree.addMutableTreeListener(listener);
    }

    public void addMutableTaxonListListener(MutableTaxonListListener listener) {
        spTree.addMutableTaxonListListener(listener);
    }

    public static Parameter createCoalPointsPopParameter(SpeciesBindings spb, Double value, Boolean bmp) {
        int dim = 0;
        for (double[] d : spb.getPopTimesSingle()) {
            dim += d.length;
        }

        if (!bmp) {
            for (double[] d : spb.getPopTimesPair()) {
                dim += d.length;
            }
        }

        return new Parameter.Default(dim, value);
    }

    public static Parameter createSplitPopulationsParameter(SpeciesBindings spb, double value, boolean root, boolean constPop) {
        int dim;
        if (constPop) {
            // one per node
            dim = 2 * spb.nSpecies() - 1;
        } else {
            // one per species leaf (ns) + 2 per internal node (2*(ns-1)) + optionally one for the root
            dim = 3 * spb.nSpecies() - 2 + (root ? 1 : 0);
        }
        return new Parameter.Default(dim, value);
    }

}
