import treemap.*;

PImage[] rhythmGroupImages;
String[][] rhythmGroupInfo;

HashMap<String, Integer> lengthCounts = new HashMap<String, Integer>();
HashMap<String, Integer> densityCounts = new HashMap<String, Integer>();
HashMap<String, Integer> typeCounts = new HashMap<String, Integer>();


int numberRGs = 4000;
String distanceMetric = "normEditDist/"; //editDist, linearity, leniency
float visibleDepthThreshold = 0.9;


RhythmGroupBinItem rootItem;
RhythmGroupItem selectedItem;
ArrayList<Integer> selectedItemRGs;
PFont font;

float treeHeight = 0;

int treemapWidth = 800;
int treemapHeight = 500;
int statusBarHeight = 200;
int scrollBarWidth = 15;
int infoBarWidth = 180;

void draw() {
  background(255);
  textFont(font);
  //draw the status bar below it
  drawStatusBar();
  
  //draw the info bar to the right
  drawInfoBar();
  
  //draw the treemap
  if (rootItem != null) {
    rootItem.draw();
  }
  
  //draw a box around the selected item
  if (selectedItem != null)
    selectedItem.drawOutline();
}


public void setup() {
  //note that this needs to be changed in the index.html file on export, too
  //does not correctly copy over variable sizes
  //currently set to 800+180 x 600+305 = 980x905
  size(treemapWidth + infoBarWidth, treemapHeight + statusBarHeight);
  rectMode(CORNERS);

  smooth();
  //noStroke();

  font = createFont("SansSerif", 13);
  parseRGTree(distanceMetric + "bins.txt", distanceMetric + "tree.txt");
  importRGInfo("rg_info.txt");
}

HashMap<Integer, String[]> binList;
HashMap<Integer, float[]> treeStructure;

void parseRGTree(String binFileName, String treeFileName) {
  //temporary storage for all the bins
  binList = new HashMap<Integer, String[]>();

  //temporary storage for tree structure
  treeStructure = new HashMap<Integer, float[]>();

  int rootId;



  //load data into the binList
  String[] lines = trim(loadStrings(binFileName));
  if (lines != null && lines.length != 0) {
    
    for (int i = 0; i < lines.length; i++) {
      println("Processing bin: " + i);
      String[] rg_ids = split(lines[i], " ");
      for (int rg_count = 0; rg_count < rg_ids.length; rg_count++) {
        rg_ids[rg_count] = rg_ids[rg_count];
      }

      //rootItem.addChild(new RhythmGroupBinItem(rootItem, "bin_" + i, i, 1, rg_ids));
      binList.put(i, rg_ids);
    }
  }

  //load data into the treeStructure
  lines = trim(loadStrings(treeFileName));
  if (lines != null && lines.length != 0) {
    //the first line contains the root node id
    rootId = int(lines[0]);
    //the remaining lists have the first number as the node and the remainder as its children
    for (int i = 1; i < lines.length; i++) {
      print(lines[i] + " -> ");
      float[] info = float(split(trim(lines[i]), " "));
      println(info);
      treeStructure.put(int(info[0]), subset(info, 1));
    }


    //now build the tree
    float initDepth = treeStructure.get(rootId)[treeStructure.get(rootId).length - 1];
    rootItem = new RhythmGroupBinItem(null, "root", 0, initDepth, null);
    for (int i = 0; i < treeStructure.get(rootId).length - 1; i++) {
      buildTree(rootItem, int(treeStructure.get(rootId)[i]), initDepth, i);
    }
  }

  if (rootItem != null) {
    rootItem.setBounds(0, 0, treemapWidth, treemapHeight);
    rootItem.contentsVisible = true;
    //rootItem.updateColors();
    assignColors();
    double numRGs = rootItem.calculateSizes();
    print("Size of the tree: " + numRGs);
    
    rhythmGroupImages = new PImage[numberRGs];
  }
}

void assignColors() {
  //first figure out how many bins there are at this depth threshold
  ArrayList<RhythmGroupItem> initColoredBins = rootItem.getBinsAtDepth(visibleDepthThreshold);
  
  //prune out parents
  ArrayList<RhythmGroupItem> removalList = new ArrayList<RhythmGroupItem>();
  for (RhythmGroupItem bin : initColoredBins) {
    //if it's a bin (i.e. if it has children)
    if (bin instanceof RhythmGroupBinItem) {
      //remove an item if one of its children is in the list as well
      for (int i = 0; i < ((RhythmGroupBinItem)bin).getItemCount(); i++) {
        RhythmGroupItem child = (RhythmGroupItem)((RhythmGroupBinItem)bin).items[i];
        if (initColoredBins.contains(child)) {
          removalList.add(bin);
          break;
        }
      }
    }
  }
  for (RhythmGroupItem bin : removalList) {
    initColoredBins.remove(bin);
  }
  
  //add in all the siblings to those bins, as they're at this threshold too!
  ArrayList<RhythmGroupItem> siblingBins = new ArrayList<RhythmGroupItem>();
  for (RhythmGroupItem bin : initColoredBins) {
    println("Found a sibling for: " + bin.id);
    siblingBins.addAll(bin.getSiblings());   
  }
  
  HashSet<RhythmGroupItem> coloredBins = new HashSet<RhythmGroupItem>();
  coloredBins.addAll(initColoredBins);
  coloredBins.addAll(siblingBins);
  
  //assign a color to each of these bins
  int index = 0;
  for (RhythmGroupItem bin : coloredBins) {
    println("Bin: " + bin.id);
    bin.hue = map(index, 0, coloredBins.size(), 0, 360);
    bin.sat = 80;
    bin.bri = 80;
    index++;
    //the children of these bins should have the same hue
    bin.propagateColorToChildren(bin.hue);
    
    //also set its parent to have visible contents
    if (bin.parent != null) {
      bin.parent.contentsVisible = true;
    }
  }
  
  assignParentColors(rootItem, coloredBins);
  
  
}

float assignParentColors(RhythmGroupBinItem item, HashSet<RhythmGroupItem> baseItems) {
  //base case: at the correct depth
  //if (item.depth <= visibleDepthThreshold) {
  if (baseItems.contains(item)) {
    return item.hue;
  }
  item.contentsVisible = true;
  //recursive case: the color is equal to the average of its children's colors
  float total = 0;
  for (int i = 0; i < item.getItemCount(); i++) {
    if (item.items[i] instanceof RhythmGroupBinItem) {
      total += assignParentColors((RhythmGroupBinItem)item.items[i], baseItems);
    }
    else if (item.items[i] instanceof RhythmGroupItem) {
      total += ((RhythmGroupItem)item.items[i]).hue;
    }
  }
  item.hue = total/item.getItemCount();
  return item.hue;
}

void keyPressed() {
  if (key == 'd') {
    distanceMetric = "leniency/";
    visibleDepthThreshold = 0.4;
    selectedItem = null;
    parseRGTree(distanceMetric + "bins.txt", distanceMetric + "tree.txt");
    
  }
}

void mousePressed() {
  boolean expandTree = false;
  if (keyPressed) {
    if (key == CODED && keyCode == SHIFT) {
      expandTree = true;
    }
  }
  if (rootItem != null && rootItem.mouseInside()) {
    rootItem.mousePressed(expandTree);
  }
  
  //DEBUG
  if (selectedItem != null) {
    println("Selected item depth: " + selectedItem.depth);
  }
}

void buildTree(RhythmGroupBinItem parent, int childID, float depth, int order) {
  println("Processing: " + parent.id + "->" + childID + " at depth: " + depth);
  if (depth > treeHeight) treeHeight = depth;
  //base case: the child is a leaf node
  if (treeStructure.get(childID) == null || childID < binList.size()) {
    parent.addChild(new RhythmGroupBinItem(parent, "bin_" + childID, order, 0, binList.get(childID)));
  }
  //otherwise, add a regular bin node
  else {
    RhythmGroupBinItem child = new RhythmGroupBinItem(parent, "node_" + childID, order, depth, null);
    parent.addChild(child);
    //and recurse for adding the child's children
    
    float[] newChildren = treeStructure.get(childID);
    float newDepth = newChildren[newChildren.length-1];
    for (int i = 0; i < newChildren.length - 1; i++) {
      buildTree(child, int(newChildren[i]), newDepth, i);
    }
  }
}

String folderContainer = "imgs/";
PImage getRGImage(int id) {
  PImage rgImage;
  if (id < 0 || id > rhythmGroupImages.length) return null;
  
  if (rhythmGroupImages[id] == null) {
    rhythmGroupImages[id] = loadImage(folderContainer + "rhythmgroup_" + id + ".png");
    //resize them to the status bar height
    int h = rhythmGroupImages[id].height;
    int w = rhythmGroupImages[id].width;
    float ratio = float(w)/float(h);
    //new width = ratio * new height
    float new_height = statusBarHeight/6;
    rhythmGroupImages[id].resize(int(ratio*new_height), int(new_height));
  }
  
  return rhythmGroupImages[id];
}

boolean imageLoaded(int id) {
  return (rhythmGroupImages[id] != null);
}

/*void importImages(String folderContainer) {
  for (int i = 0; i < rhythmGroupImages.length; i++)
  {
    rhythmGroupImages[i] = loadImage(folderContainer + "rhythmgroup_" + i + ".png");
    //resize them to the status bar height
    int h = rhythmGroupImages[i].height;
    int w = rhythmGroupImages[i].width;
    float ratio = float(w)/float(h);
    //new width = ratio * new height
    float new_height = statusBarHeight/6;
    rhythmGroupImages[i].resize(int(ratio*new_height), int(new_height));
  }
}*/

void importRGInfo(String filename) {
  //initialize the counter hashes
  lengthCounts.put("5", 0);  lengthCounts.put("10", 0);  lengthCounts.put("15", 0);  lengthCounts.put("20", 0);
  densityCounts.put("low", 0);  densityCounts.put("medium", 0);  densityCounts.put("high", 0);
  typeCounts.put("swing", 0);  typeCounts.put("regular", 0);  typeCounts.put("random", 0);
  
  String[] info = loadStrings(filename);
  rhythmGroupInfo = new String[info.length][3];
  for (int i = 0; i < info.length; i++) {
    rhythmGroupInfo[i] = split(trim(info[i]), " ");
    //length, type, density
    String rg_length = rhythmGroupInfo[i][0];
    String rg_type = rhythmGroupInfo[i][1];
    String rg_density = rhythmGroupInfo[i][2];
    lengthCounts.put(rg_length, lengthCounts.get(rg_length) + 1);
    typeCounts.put(rg_type, typeCounts.get(rg_type) + 1);
    densityCounts.put(rg_density, densityCounts.get(rg_density) + 1);
  }
}