//Copyright 2005 Sean McCullough
//banksean at yahoo

int H = 600;
int W = 800;
float G = 0.750;

float scaleFactor = 1.0;
float transX = 0;
float transY = 0;

HashMap gdpByNation = new HashMap();
Graph g = new Graph();
 

PFont font; 
// The font must be located in the sketch's 
// "data" directory to load successfully 
boolean doneLoading = false;
int fCount = 0;

void setup() {
  size(W,H);
  framerate(30);
  smooth();
  font = loadFont("EurostileRegular-14.vlw"); 
  textFont(font, 14); 
  g = buildExportAtlasGraph();
  doneLoading= true;
}

void draw() {
  background(0);
  translate(transX, transY);
  scale((float)scaleFactor);
  if (g != null && doneLoading) {
    doLayout();
    if (fCount > 100) {
      g.draw();  
    }
    fCount++;
  }
}

void doLayout() {
  if (g == null || !doneLoading) {
     return;
  }
  //calculate forces on each node
  //calculate spring forces on each node
  for (int i=0; i<g.getNodes().size(); i++) {
    ForcedNode n = (ForcedNode)g.getNodes().get(i);
    ArrayList edges = (ArrayList)g.getEdgesFrom(n);
    n.setForce(new Vector3D(0,0,0));
    for (int j=0; edges != null && j<edges.size(); j++) {
      SpringEdge e = (SpringEdge)edges.get(j);
      Vector3D f = e.getForceFrom();
      n.applyForce(f);
    }
    
    edges = (ArrayList)g.getEdgesTo(n);
    for (int j=0; edges != null && j<edges.size(); j++) {
      SpringEdge e = (SpringEdge)edges.get(j);
      Vector3D f = e.getForceTo();
      n.applyForce(f);
    }
    
  }
  
  //calculate the anti-gravitational forces on each node
  //this is the N^2 shittiness that needs to be optimized
  //TODO: at least make it N^2/2 since forces are symmetrical
  for (int i=0; i<g.getNodes().size(); i++) {
    ForcedNode a = (ForcedNode)g.getNodes().get(i);
    for (int j=0; j<g.getNodes().size(); j++) {
      ForcedNode b = (ForcedNode)g.getNodes().get(j);
      if (b != a) {
        float dx = b.getX() - a.getX();
        float dy = b.getY() - a.getY();
        float r = (float)sqrt((float)(dx*dx + dy*dy));
        //F = G*m1*m2/r^2  
        
        if (r != 0) { //don't divide by zero.
          float f = G*(a.getMass()*b.getMass()/(r*r));
          Vector3D vf = new Vector3D(-dx*f, -dy*f, 0);
          a.applyForce(vf);
        }              
      }
    }
  }
  
  //move nodes according to forces
  for (int i=0; i<g.getNodes().size(); i++) {
    ForcedNode n = (ForcedNode)g.getNodes().get(i);
    if (n != g.getDragNode()) {
      n.setPosition(n.getPosition().add(n.getForce()));
    }
  }
}

void keyPressed() {
  if (key == '[') {
    scaleFactor += 0.1;
  } else if (key == ']') {
    scaleFactor -= 0.1;
  }
}

void mousePressed() {
  //g.setSelectedNode(null);
  //g.setDragNode(null);
  for(int i=0; i<g.getNodes().size(); i++) {
    Node n = (Node)g.getNodes().get(i);
    if (n.containsPoint((mouseX-transX)/scaleFactor, (mouseY-transY)/scaleFactor)) {
      g.setSelectedNode(n);
      g.setDragNode(n);
    }
  }
}

void mouseMoved() {
  if (g.getDragNode() == null) {
    g.setHoverNode(null);
    for(int i=0; i<g.getNodes().size(); i++) {
      Node n = (Node)g.getNodes().get(i);
      if (n.containsPoint((mouseX-transX)/scaleFactor, (mouseY-transY)/scaleFactor)) {
        g.setHoverNode(n);
      }
    }
  }
}

void mouseReleased() {
  g.setDragNode(null);
}

void mouseDragged() {
  if (g.getDragNode() != null) {
    g.getDragNode().setPosition(new Vector3D((mouseX-transX)/scaleFactor, (mouseY-transY)/scaleFactor, 0));
  } else {
    transX += mouseX-pmouseX;
    transY += mouseY-pmouseY;
  }
}

String normalizeNation(String n) {
    if (n.equals("US")) {
      return "United States";
    }
    if (n.equals("UK")) {
      return "United Kingdom";
    }
    if (n.equals("UAE")) {
      return "United Arab Emirates";
    }
    if (n.equals("North Korea")) {
      return "Korea, North";
    }
    if (n.equals("South Korea")) {
      return "Korea, South";
    }
    if (n.equals("EU")) {
      return "European Union";
    }
    
    return n;
}

Graph buildExportAtlasGraph() {
  int maxCountries = 3000;
  
  //First get the GDP, then load the export graph.
  String gdps[] = loadStrings("gdp.txt");
  for (int i = 0; i<gdps.length && i<maxCountries; i++) {
    if (gdps[i].indexOf('\t') != -1) {
      String nation = normalizeNation(gdps[i].substring(0,gdps[i].indexOf('\t')).trim());
      String gdpStr = gdps[i].substring(nation.length()).trim().substring(1);
      if (gdpStr.indexOf(" ") != -1) {
        String gdpN = gdpStr.substring(0, gdpStr.indexOf(" "));
        try {
          float gdp = Float.parseFloat(gdpN);
          if (gdpStr.indexOf("million") != -1) {
            gdp = gdp*1;
            println(nation + ": " + Math.log(gdp));
            //gdpByNation.put(nation, new Float(gdp));
          }
          if (gdpStr.indexOf("billion") != -1) {
            gdp = gdp;
            println(nation + ": " + Math.log(gdp));
            //gdpByNation.put(nation, new Float(gdp));
          }
          if (gdpStr.indexOf("trillion") != -1) {
            //gdp = gdp*1000;
            println(nation + ": " + gdp);
            gdpByNation.put(nation, new Float(gdp));
          }
        } catch (Exception ex) {
          println("couldn't parse GDP for " + nation + ": " + gdpStr);
        }
      }
    }
  }
  
  String lines[] = loadStrings("exports.txt"); 

  for (int i = 0; i<lines.length && i<maxCountries; i++) {
    String exportNation = normalizeNation(lines[i].substring(0,lines[i].indexOf('\t')).trim());
    String importNations = lines[i].substring(exportNation.length()).trim();
    importNations = importNations.substring(0,importNations.length()-6);
    StringTokenizer st = new StringTokenizer(importNations, "%");
    println(exportNation + ": ");
    ForcedNode exNode = getNationNode(exportNation);
    
    Float egdp = (Float)gdpByNation.get(exportNation);

      while (egdp != null && st.hasMoreTokens()) {
      String importNation = st.nextToken().trim();
      if (importNation.length() > 0) {
        if (importNation.startsWith(", ")) {
           importNation = importNation.substring(2);
        }
        if (importNation.lastIndexOf(" ") != -1) {
          String percent = importNation.substring(importNation.lastIndexOf(" "));
          importNation = normalizeNation(importNation.substring(0,importNation.lastIndexOf(" ")));
          
          try {
            float pct = Float.parseFloat(percent);
            println("\t" + importNation + ": " + pct);

            ForcedNode imNode = getNationNode(importNation);
            Float gdp = (Float)gdpByNation.get(importNation);
            
              if (gdp != null) {
                SpringEdge edge = new SpringEdge(imNode, exNode);
                edge.setNaturalLength(250-(float)pct);
                edge.setStiffness(0.001*pct);  //0.05);
                edge.setWeight(pct/10);
                edge.setLabel(pct + "%");
                g.addEdge(edge);
              } else {
                println("couldn't find GDP for " + importNation);
                imNode.setMass(1);
              } 
          } catch (Exception ex) {
            ex.printStackTrace();
          }
          
        } else {
          println ("\t " + importNation + " no percent"); 
        }
      }
    }
  }
  return g;
}

ForcedNode getNationNode(String name) {
    ForcedNode exNode = (ForcedNode)g.getNode(name);
    if (exNode == null) {
      Float egdp = (Float)gdpByNation.get(name);
      exNode = new ForcedNode(new Vector3D(W/4 + random(W/2), H/4 + random(H/2), 0));
      exNode.setId(name);
      exNode.setLabel(name);
      if (egdp != null) {
        exNode.setMass((float)((Math.sqrt(egdp.floatValue()))));
        exNode.setAltLabel("$" + egdp);
        g.addNode(exNode);
      } else {
        println("couldn't find GDP for " + name);
        exNode.setMass(1);
      }
    }
    return exNode;
}