package dr.app.treespace;

import dr.xml.Reference;
import jebl.evolution.graphs.Node;
import jebl.evolution.trees.RootedTree;

import java.awt.geom.Rectangle2D;
import java.util.*;

/**
 * @author Andrew Rambaut
 * @version $Id$
 */
public class TreeLineages {

    private final static String[] LOCATIONS = {
            "Africa",
            "USA",
            "Taiwan",
            "China",
            "Russia",
            "Oceania",
            "Asia",
            "Japan",
            "Mexico",
            "South America",
            "Canada",
            "Europe",
            "Southeast Asia",
            "South Korea"
    };

    private final static String[] AIR_COMMUNITIES = {
            "AC1",
            "AC2",
            "AC3",
            "AC4",
            "AC5",
            "AC6",
            "AC7",
            "AC8",
            "AC9",
            "AC10",
            "AC11",
            "AC12",
            "AC13",
            "AC14"
    };

    Map<String, Integer> locationMap = new HashMap<String, Integer>();
    private static final double THRESHOLD = 0.01;

    public TreeLineages() {
        for (int i = 0; i < AIR_COMMUNITIES.length; i++) {
            locationMap.put(AIR_COMMUNITIES[i], i);
        }
    }

    public List<Lineage> getRootLineages() {
        return rootLineages;
    }

    public Rectangle2D getTreeBounds() {
        return treeBounds;
    }

    public Rectangle2D getAntigenicBounds() {
        return antigenicBounds;
    }

    public void addTree(RootedTree tree) {

        // the offset is the distance to the right most tip - used to align the tips of the tree
        offsetX = 0.0;
        Lineage lineage = new Lineage();
        addNode(tree, tree.getRootNode(), lineage, 0.0);

        lineage.dx = offsetX;

        rootLineages.add(lineage);
    }

    public void setupTrees() {
        for (BitSet key : nodeCounts.keySet()) {
            int count = nodeCounts.get(key);
            if ((((double)count) / rootLineages.size()) > THRESHOLD) {
                majorityNodes.add(key);
            }
        }
        Collections.sort(majorityNodes, new Comparator<BitSet>() {
            public int compare(BitSet bitSet, BitSet bitSet1) {
                return bitSet.cardinality() - bitSet1.cardinality();
            }
        });

        for (Lineage rootLineage : rootLineages) {
        // set the tip count to zero for this traversal
//        currentY = 0.0;
            positionNode(rootLineage);
        }
    }

    public BitSet addNode(RootedTree tree, Node node, Lineage lineage, double cumulativeX) {
        lineage.bits = new BitSet();

        cumulativeX += tree.getLength(node);
        lineage.dx = cumulativeX;

        String location = (String)node.getAttribute("states");
        if (location != null) {
            lineage.state = locationMap.get(location);
        }

        if (!tree.isExternal(node)) {
            List<Node> children = tree.getChildren(node);
            if (children.size() != 2) {
                throw new RuntimeException("Tree is not binary");
            }

            lineage.child1 = new Lineage();
            lineage.child2 = new Lineage();

            lineage.bits.or(addNode(tree, children.get(0), lineage.child1, cumulativeX));
            lineage.bits.or(addNode(tree, children.get(1), lineage.child2, cumulativeX));

            lineage.tipCount = lineage.child1.tipCount + lineage.child2.tipCount;

            if (lineage.child1.tipCount > lineage.child2.tipCount ||
                    (lineage.child1.tipCount == lineage.child2.tipCount &&
                            lineage.child1.tipNumber > lineage.child2.tipNumber)) {
                Lineage tmp = lineage.child1;
                lineage.child1 = lineage.child2;
                lineage.child2 = tmp;
            }

            Integer count = nodeCounts.get(lineage.bits);
            if (count == null) {
                count = 1;
            } else {
                count ++;
            }
            nodeCounts.put(lineage.bits, count);

        } else {
            Integer tipNumber = taxonNumbers.get(tree.getTaxon(node).getName());
            if (tipNumber == null) {
                tipNumber = taxonNumbers.size();
                taxonNumbers.put(tree.getTaxon(node).getName(), tipNumber);
            }

            lineage.tipNumber = tipNumber;
            lineage.tipCount = 1;

            lineage.bits.set(lineage.tipNumber);

            if (cumulativeX > offsetX) {
                offsetX = cumulativeX;
            }
        }

        Object[] ag = (Object[])node.getAttribute("antigenic");
        if (ag != null) {
//            lineage.ax = (Double)ag[0];
//            lineage.ay = (Double)ag[1];
            lineage.ax = lineage.dx;
            lineage.ay = (Double)ag[0];

            antigenicBounds.add(lineage.ax, lineage.ay);
        } else {
            Double ag1 = (Double)node.getAttribute("antigenic1");
            Double ag2 = (Double)node.getAttribute("antigenic2");
            if (ag1 != null) {
                double ad =

                lineage.ay = ag1;
//                lineage.ay = ag2;
                lineage.ax = lineage.dx;

                antigenicBounds.add(lineage.ax, lineage.ay);
            }

        }

        return lineage.bits;
    }

    public double positionNode(Lineage lineage) {
        if (lineage.child1 != null) {
            lineage.dy = positionNode(lineage.child1);
            lineage.dy += positionNode(lineage.child2);

            // the y of this node is the average of the two children
            lineage.dy /= 2.0;

            // now change the children to relative y positions.
//            lineage.child1.dy = lineage.child1.dy - lineage.dy;
//            lineage.child2.dy = lineage.child2.dy - lineage.dy;

            Double location = nodeLocations.get(lineage.bits);
            if (location == null) {
                BitSet bestBits = lineage.bits;

                if (!majorityNodes.contains(bestBits)) {
                    // otherwise find a bigger node...
                    for (BitSet bits : majorityNodes) {
                        BitSet bits1 = (BitSet)bits.clone();
                        bits1.and(lineage.bits);
                        if (bits1.cardinality() == lineage.bits.cardinality()) {
                            bestBits = bits;
                            break;
                        }
                    }

                    location = nodeLocations.get(bestBits);
                    if (location != null) {
                        lineage.dy = location;
                    }
                }

                nodeLocations.put(bestBits, lineage.dy);

            } else {
                lineage.dy = location;
            }

        } else {
            Integer orderedTipNumber = orderedTipNumbers.get(lineage.tipNumber);
            if (orderedTipNumber == null) {
                orderedTipNumber = orderedTipNumbers.size();
                orderedTipNumbers.put(lineage.tipNumber, orderedTipNumber);
            }

            lineage.tipNumber = orderedTipNumber;
//            lineage.dy = LATITUDES[lineage.state];
            lineage.dy = lineage.tipNumber;
//            lineage.dy = currentY;
//            currentY += 1.0;
        }

        treeBounds.add(lineage.dx, lineage.dy);

        return lineage.dy;
    }

    class Lineage {
        double dx = 0;
        double dy = 0;
        double ax = 0;
        double ay = 0;
        Lineage child1 = null;
        Lineage child2 = null;
        int tipNumber = 0;
        int tipCount = 0;
        int state;
        BitSet bits;
    }

    private List<Lineage> rootLineages = new ArrayList<Lineage>();
    private Map<String, Integer> taxonNumbers = new HashMap<String, Integer>();
    private Map<Integer, Integer> orderedTipNumbers = new HashMap<Integer, Integer>();
    private Map<BitSet, Integer> nodeCounts = new HashMap<BitSet, Integer>();
    private Map<BitSet, Double> nodeLocations = new HashMap<BitSet, Double>();
    private List<BitSet> majorityNodes = new ArrayList<BitSet>();

    private double offsetX;

    private Rectangle2D treeBounds = new Rectangle2D.Double();
    private Rectangle2D antigenicBounds = new Rectangle2D.Double();
}
