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.inference.model.*;
import dr.inference.operators.OperatorFailedException;
import dr.inference.operators.Scalable;
import dr.util.Attributable;
import dr.util.HeapSort;
import dr.xml.*;
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, NodeAttributeProvider, TreeLogger.LogUpon, Scalable {
    public static final String SPECIES_TREE = "speciesTree";

    public static final String SPP_SPLIT_POPULATIONS = "sppSplitPopulations";
    public static final String COALESCENT_POINTS_POPULATIONS = "coalescentPointsPopulations";
    public static final String COALESCENT_POINTS_INDICATORS = "coalescentPointsIndicators";

    public static final String BMPRIOR = "bmPrior";
    public static final String CONST_ROOT_POPULATION = "constantRoot";
    public static final String CONSTANT_POPULATION = "constantPopulation";

    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());
        }
    }

    SpeciesTreeModel(SpeciesBindings species, Parameter sppSplitPopulations,
                     Parameter coalPointsPops, Parameter coalPointsIndicator, Tree startTree,
                     boolean bmp, boolean nonConstRootPopulation, boolean constantPopulation) {
        super(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);

        while( check ) {
            // Start tree had branch info, keep it when compatible, decrease height by 1% until
            // compatible otherwise.
            check = false;
            for( SpeciesBindings.GeneTreeInfo t : species.getGeneTrees() ) {
                if( ! isCompatible(t) ) {
                    SimpleTree.Utils.scaleNodeHeights(spTree, 0.99);
                    check = true;
                    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( 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 p0;
    }

    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(int state) {
        final String curTop = Tree.Utils.uniqueNewick(spTree, spTree.getRoot());
        if( state == 0 || !curTop.equals(previousTopology) ) {
            previousTopology = curTop;
            return true;
        }
        return false;
    }

    public String[] getNodeAttributeLabel() {
        // keep short, repeated endlessly in tree log
        return new String[]{"dmf"};
    }

    public String[] getAttributeForNode(Tree tree, NodeRef node) {
        assert tree == this;

        //final VDdemographicFunction df = getProps().get(node).demogf;

        final DemographicFunction df = getNodeDemographic(node);
        return new String[]{"{" + df.toString() + "}"};
    }

    // 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) {
        return spTree.getTaxonIndex(taxon);
    }

    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);
    }

    private 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);
    }

    private 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);
    }

    public static XMLObjectParser PARSER = new AbstractXMLObjectParser() {

        public Object parseXMLObject(XMLObject xo) throws XMLParseException {
            SpeciesBindings spb = (SpeciesBindings) xo.getChild(SpeciesBindings.class);

            Parameter coalPointsPops = null;
            Parameter coalPointsIndicators = null;
            final Boolean cr = xo.getAttribute(CONST_ROOT_POPULATION, false);
            final Boolean cp = xo.getAttribute(CONSTANT_POPULATION, false);

            final Boolean bmp = xo.getAttribute(BMPRIOR, false);
            {
                XMLObject cxo = xo.getChild(COALESCENT_POINTS_POPULATIONS);
                if( cxo != null ) {
                    final double value = cxo.getAttribute(Attributable.VALUE, 1.0);
                    coalPointsPops = createCoalPointsPopParameter(spb, cxo.getAttribute(Attributable.VALUE, value), bmp);
                    ParameterParser.replaceParameter(cxo, coalPointsPops);
                    coalPointsPops.addBounds(
                            new Parameter.DefaultBounds(Double.MAX_VALUE, 0, coalPointsPops.getDimension()));

                    cxo = xo.getChild(COALESCENT_POINTS_INDICATORS);
                    if( cxo == null ) {
                        throw new XMLParseException("Must have indicators");
                    }
                    coalPointsIndicators = new Parameter.Default(coalPointsPops.getDimension(), 0);
                    ParameterParser.replaceParameter(cxo, coalPointsIndicators);
                } else {
                   // assert ! bmp;
                }
            }

            final XMLObject cxo = xo.getChild(SPP_SPLIT_POPULATIONS);

            final double value = cxo.getAttribute(Attributable.VALUE, 1.0);
            final boolean nonConstRootPopulation = coalPointsPops == null && !cr;
            final Parameter sppSplitPopulations = createSplitPopulationsParameter(spb, value, nonConstRootPopulation, cp);
            ParameterParser.replaceParameter(cxo, sppSplitPopulations);

            final Parameter.DefaultBounds bounds =
                    new Parameter.DefaultBounds(Double.MAX_VALUE, 0, sppSplitPopulations.getDimension());
            sppSplitPopulations.addBounds(bounds);

            final Tree startTree = (Tree) xo.getChild(Tree.class);

            return new SpeciesTreeModel(spb, sppSplitPopulations, coalPointsPops, coalPointsIndicators, startTree, bmp,
                    nonConstRootPopulation, cp);
        }

        public XMLSyntaxRule[] getSyntaxRules() {
            return new XMLSyntaxRule[]{
                    AttributeRule.newBooleanRule(BMPRIOR, true),
                    AttributeRule.newBooleanRule(CONST_ROOT_POPULATION, true),
                     AttributeRule.newBooleanRule(CONSTANT_POPULATION, true),
                    new ElementRule(SpeciesBindings.class),
                    // A starting tree. Can be very minimal, i.e. no branch lengths and not resolved
                    new ElementRule(Tree.class, true),
                    new ElementRule(SPP_SPLIT_POPULATIONS, new XMLSyntaxRule[]{
                            AttributeRule.newDoubleRule(Attributable.VALUE, true),
                            new ElementRule(Parameter.class)}),

                    new ElementRule(COALESCENT_POINTS_POPULATIONS, new XMLSyntaxRule[]{
                            AttributeRule.newDoubleRule(Attributable.VALUE, true),
                            new ElementRule(Parameter.class)}, true),

                    new ElementRule(COALESCENT_POINTS_INDICATORS, new XMLSyntaxRule[]{
                            new ElementRule(Parameter.class)}, true),
            };
        }

        public String getParserDescription() {
            return "Species tree which includes demographic function per branch.";
        }

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

        public String getParserName() {
            return SPECIES_TREE;
        }
    };
}
