//----------------------------------------------------------------------------

//----------------------------------------------------------------------------


//----------------------------------------------------------------------------
// User-modifiable variables
int maxBranches = 200;
int maxBifurs = 2;  // maximum number of splits at each branch terminus
float maxBifurA = 55.0;
float minBifurA = 20.0;
float pBranch = 0.4;
float pExt    = 0.02;
float pRet    = 0.55;
int minExt = 8; 
int maxExt = 22;
int minRetract = 5;
int maxRetract = 30;
int delThresh = 4;
int endField = 5;  // distance from end which can sense mouse
float incSize = 0.05;
int histSlots = 40;
int fade = 35;
int inhibLeak = -1;  // make proportional to inhib value?
//----------------------------------------------------------------------------

//----------------------------------------------------------------------------
// Boring variables
Neuron n;
Branch[] Bs = new Branch[maxBranches];
int numBranches = 0;
int maxID = 0;
int overB = -1;
int maxDistance = 1;

int attSize = 50;
int attW = 12;
int attH = 10;
int[][] att = new int[attW][attH];  // attractant map
boolean showAttMap = true;

int inhibSize = 25;
int inhibW = 24;
int inhibH = 20;
int[][] inhib = new int[inhibW][inhibH];  // attractant map
boolean showInhibMap = true;

String mode = "stopped";
String adjMode = "pBranch";

PFont norm;
PFont bold;
float simVar;

int[] branchHist = new int[40];
boolean histOn = false;
boolean statOn = false;
boolean printProbs = true;
int colorRange = 0;
int numColorMaps = 3;

boolean shiftDown = false;
boolean controlDown = false;
boolean altDown = false;
//----------------------------------------------------------------------------


//----------------------------------------------------------------------------
void setup()
{
  size(600, 500);
  frameRate(600);
  n = new Neuron(width/2, height-50, 30, 270, 0);
  Bs[0] = new Branch(n.dendrootX, n.dendrootY, n.rootAngle, 25.0, 0, 0, 0); // proximal dendrite - make random?
  numBranches++;

  norm = loadFont("BankGothic-Light-12.vlw");
  bold = loadFont("BankGothic-Medium-12.vlw");
  textFont(norm);

  // initialize maps...
  for(int i=0; i<attH; i++){
    for(int j=0; j<attW; j++){
      att[j][i] = (attH-i)*10+40;
    }
  }

  for(int i=0; i<inhibH; i++){
    for(int j=0; j<inhibW; j++){
      inhib[j][i] = (inhibW-j)*3+10;
    }
  }

  background(151);
}
//----------------------------------------------------------------------------


//----------------------------------------------------------------------------
// loopage
void draw()
{
  fill(0, fade);
  noStroke();
  rect(0, 0, width, height);

  // check if the mouse is hovering over anything
  //overB = overAny();

  if(frameCount%60 == 0){
    println(frameRate);
  }

  // draw attractant map...
  if(showAttMap){
    noStroke();
    for(int i=0; i<attH; i++){
      for(int j=0; j<attW; j++){
        int a = att[j][i];
        if(a > 0){
          fill(a, a, a, fade);
          rect(j*attSize, i*attSize, attSize, attSize);
        }
      }
    }
  }

  // draw inhibitory map...
  noStroke();
  for(int i=0; i<inhibH; i++){
    for(int j=0; j<inhibW; j++){
      int iVal = inhib[j][i];
      if(frameCount%2==0 && iVal>15){
        inhib[j][i] += inhibLeak;      // leakage
      }
      if(showInhibMap){   
        if(iVal > 0){
          fill(iVal/4, iVal, 0, fade/3 + 2);
          rect(j*inhibSize, i*inhibSize, inhibSize, inhibSize);
        }
      }
    }
  }

  // draw soma...
  n.drawNeuron(); 

  // internal loop for any functions which operate on each branch in the tree...
  for(int i=0; i<numBranches; i++){
    // check if mouse is over any branches
    overB = Bs[i].over();

    // draw branch...
    if(colorRange == 0){ // color based on age on initial branch
      int c = int(((1.0*Bs[i].age)/Bs[0].age)*255);
      stroke(180, c, c, 50+Bs[i].voltage*2);
    } 
    else if(colorRange == 1){
      stroke(180, Bs[i].age/5, Bs[i].age/5, 50+Bs[i].voltage*2);
    }
    else if(colorRange == 2){
      int c = int((Bs[i].distance/float(maxDistance))*255);  // color normalized to branch distance from soma (based on longest ever branch length)
      stroke(180, c, c, 50+Bs[i].voltage*2);
    } 
    else {
      stroke(50, 200, 60, 50+Bs[i].voltage*2);
    }
    strokeWeight(3);
    if(frameCount%1==0){
      //line(Bs[i].startX, Bs[i].startY, Bs[i].endX, Bs[i].endY);
      Bs[i].drawBranch();
    }
    
    // leak voltage...
    Bs[i].voltage = int(Bs[i].voltage*0.95 + Bs[Bs[i].parent].voltage*0.25);

    // get action probabilites based on environment...
    int x = Bs[i].endX;
    int y = Bs[i].endY;
    if(x<width && x>0 && y<height && y>0){  // if branch within screen
      // environment...
      int attX = int(Bs[i].endX/attSize);
      int attY = int(Bs[i].endY/attSize);
      int a = att[attX][attY];

      int inhibX = int(x/inhibSize);  // these three must be outside next if statement so they are availbe for the growth if statement
      int inhibY = int(y/inhibSize);
      int in = inhib[inhibX][inhibY];

      // adjust probabilities for this branch...
      // branching probability
      if(Bs[i].age > 2000 && !Bs[i].filo){
        pBranch = a/1025.0;
      }
      else {
        pBranch = a/255.0;
      }

      // retraction probability
      pRet = in/(3235*log(Bs[i].age+900)-21750);

      // extension probability
      if(Bs[i].filo && Bs[i].len<maxExt*2){
        pExt = 0.03;
      }
      else {
        pExt = 0.0001;
      }

      // if mouse is over a branch...
      if(overB==2){
        if(mousePressed){
          Bs[i].voltage +=10;    // increase voltage
          println(Bs[i].voltage);
        }
        
        // draw text...
        if(printProbs){
          fill(0);
          noStroke();
          rect(0, 0, 130, 100);
          fill(255, 180);
          //text(mode, 5, 10);
          text("over branch # " + overB, 10, 10);
          text("probabilities of...", 5, 24);
          text("branching: " + nf(pBranch, 1, 2), 10, 38);
          text("extension: " + nf(pExt, 1, 2), 10, 52);
          text("retraction: " + nf(pRet, 1, 2), 10, 66);
          text("age: " + str(Bs[i].age), 10, 80); 
          text("distance: " + str(Bs[i].distance), 10, 90);
        }
        Bs[i].highlight();
      }


      // expand dendritic arbor...
      if(mode == "growing"){
        Bs[i].age++;

        // lay down some inhibitory juice if the branch is stable 
        if(Bs[i].age>120 && in<255){
          inhib[inhibX][inhibY] = in + 1;
        }

        // select growth action...
        float r = random(1);
        if(r <= pBranch && numBranches<maxBranches && Bs[i].bifurs<maxBifurs){  // bifurcate if there are branch slots available
          Bs[numBranches] = Bs[i].bifurcate();
        }
        else if(r > pBranch && r <= pBranch+pExt && Bs[i].filo){
          Bs[i].extend();
        }
        else if (r > pBranch+pExt && r <= pBranch+pExt+pRet && Bs[i].filo && numBranches>1){
          Bs[i].retract(); 
        }
      }
      else {  // if not growing...

      }
    }
    else {    // if branch is outside screen boundaries
      Bs[i].retract();
    }
  }


  // modify maps...
  if(mousePressed){
    if(controlDown){
      int attX = mouseX/attSize;
      int attY = mouseY/attSize;
      if(shiftDown){
        att[attX][attY] -= 1;
        println(att[attX][attY]);
      }
      else {
        att[attX][attY]++;
        println(att[attX][attY]);
      }
    } 
    else {  // if control is not pressed, modify inhibatory map
      int inhibX = mouseX/inhibSize;
      int inhibY = mouseY/inhibSize;
      if(shiftDown){
        inhib[inhibX][inhibY] -= 4;
        println(inhib[inhibX][inhibY]);
      } 
      else {
        inhib[inhibX][inhibY] += 4;
        //println(inhib[inhibX][inhibY]);
      }
    }
  }


  // print stats...
  if(statOn){
    if(frameCount%12 == 0){
      stroke(200);
      stats();
    }
  }


  // recalc histogram...
  if(histOn || (mouseX<30 && mouseY>height-30)){
    calcHist();
    for(int i=0; i<branchHist.length; i++){
      line(6*i+30, height-30, 6*i+30, branchHist[i]*-3 + height-30);
    }
  }


  // change fade if mouse on right side...
  if(mouseX>width-20){
    fade = int((mouseY*55)/height);
    //println("fade: " + fade);
    line(width-20, mouseY, width, mouseY);
  }
}


//----------------------------------------------------------------------------
void mousePressed(){
  if(mouseX<30 && mouseY>height-30){  // keep histogram on if corner clicked
    if(histOn){
      histOn = false;
    } 
    else {
      histOn = true;
    }
  }
}


//----------------------------------------------------------------------------
void keyPressed()
{
  if(key == 'g'){
    if(mode == "growing"){
      mode = "stopped";
    } 
    else{
      mode = "growing";
      println("input mode: " + mode);
    }
  }

  if(key == 'm'){
    if(showAttMap){
      showAttMap = false;
    } 
    else {
      showAttMap = true;
    }
  }

  if(key == 'n'){
    if(showInhibMap){
      showInhibMap = false;
    } 
    else {
      showInhibMap = true;
    }
  }

  if(key == 'c'){
    colorRange = (colorRange+1)%numColorMaps;
    if(colorRange == 0){
      println("Branch color normalized to initial branch age");
    } 
    else if(colorRange == 1){
      println("Branch color set to age, saturates at about 1000 timesteps");
    }
    else if(colorRange == 2){
      println("Branch color proportional to distance from soma"); 
    }
    else{
      println("----- error selecting color mapping: colorRange = " + colorRange + " -----");
    }
  }

  if(key == 'p'){
    if(!printProbs){
      printProbs = true;
      println("print branch probability on");
    } 
    else {
      printProbs = false;
      println("print branch probability off");
    }
  }

  if(key == 'h'){
    if(!histOn){
      histOn = true;
      println("Histogram display on");
    } 
    else {
      histOn = false;
      println("Histogram display off");
    }
  }

  if(key == 't'){
    if(!statOn){
      statOn = true;
      println("Statistics display on");
    } 
    else {
      statOn = false;
      println("Statistics display off");
    }
  }

  if(key == 'b'){
    if(numBranches < maxBranches){
      int bifur = int(random(0, numBranches-1));
      if(Bs[bifur].bifurs < maxBifurs){
        println("new branch at " + bifur);
        Bs[numBranches] = Bs[bifur].bifurcate();
      }
    } 
    else {
      println("encountered branch limit");
    }
  }

  if(key == 'e'){
    for(int i=0; i<numBranches; i++){
      if(Bs[i].filo){
        Bs[i].extend();
      }
    }
  }

  if(key == 'r'){
    if(numBranches>1){
      for(int i=0; i<numBranches; i++){
        if(Bs[i].filo){
          Bs[i].retract();
        }
      }
    }
  }

  if(key == 'x'){  // reset
    Branch[] newBs = new Branch[maxBranches];
    numBranches = 0;
    newBs[0] = new Branch(n.dendrootX, n.dendrootY, n.rootAngle, 25.0, 0, 0); // proximal dendrite - make random?
    newBs[0].bifurs = 0;
    newBs[0].filo = true;
    Bs = newBs;
    numBranches = 1;
  }

  if(key =='t'){
    stats();
  }

  if(key == 'a'){
    adjMode = "pBranch";
  }

  if(key == 's'){
    println("snapshot taken at " + frameCount);
    saveFrame("neuron-###.png");    // take a snap shot if 's' button pressed
  }

  if(key == 'd'){
    adjMode = "pRet";
  }

  // increment var up a small amount...
  if(keyCode == UP){
    numVarInc(incSize);
  }

  // increment var down a small amount...
  if(keyCode == DOWN){
    numVarInc(-1*incSize);
  }

  //---------------------------------------------
  // modifier mode modifiers...
  if(keyCode == SHIFT){
    shiftDown = true;
  }

  if(keyCode == CONTROL){
    controlDown = true;
  }

  if(keyCode == ALT){
    altDown = true;
  }
}

//----------------------------------------------------------------------------
void keyReleased(){
  if(keyCode == SHIFT){
    shiftDown = false;
  }

  if(keyCode == CONTROL){
    controlDown = false;
  }

  if(keyCode == ALT){
    altDown = false;
  }
}



//----------------------------------------------------------------------------
//------ HELPERS -------------------------------------------------------------
//----------------------------------------------------------------------------

//----------------------------------------------------------------------------
void stats(){
  float cumLen = 0.0;
  int cumOrder = 0;
  int cumBifurs = 0;
  int maxOrder = 0;
  float cumAngle = 0.0;
  int numFilos = 0;
  int maxX = 0;
  int maxY = 0;
  int minX = 1000;
  int minY = 1000;

  for(int i=0; i<numBranches; i++){
    cumLen = cumLen + Bs[i].len;
    int curO = Bs[i].order;

    // find rough area...
    if(Bs[i].endX > maxX){
      maxX = Bs[i].endX;
    } 
    else if(Bs[i].endX < minX){
      minX = Bs[i].endX;
    }
    if(Bs[i].endY > maxY){
      maxY = Bs[i].endY;
    } 
    else if(Bs[i].endY < minY){
      minY = Bs[i].endY;
    }

    // find max order...
    if(maxOrder < curO){
      maxOrder = curO;
    }

    // update cumulative variables... 
    cumBifurs = cumBifurs + Bs[i].bifurs;
    cumAngle = cumAngle + Bs[i].angle;
    cumOrder = cumOrder + curO;
    if(Bs[i].filo){
      numFilos++;
    }
  }
  float avgLen = cumLen/numBranches;
  float avgOrder = cumOrder/(float)numBranches;
  float avgBifurs = cumBifurs/(float)numBranches;
  float avgAngle = cumAngle/numBranches;
  int maxWidth = maxX-minX;
  int maxHeight = maxY-minY;
  int area = maxWidth*maxHeight;

  println("Stats: average length: " + nf(avgLen, 2, 1) + " | total length: " + cumLen 
    + " | average order: " + nf(avgOrder, 1, 1) + " | max order: " + maxOrder 
    + " | average bifurcations: " + nf(avgBifurs, 1, 2) 
    + " | average angle: " + nf(avgAngle, 1, 2) + " | num filopodia: " + numFilos
    + " | rough px area: " + area);
}


//----------------------------------------------------------------------------
void calcHist(){
  branchHist = new int[histSlots];
  for(int i=0; i<numBranches; i++){
    int curO = int(Bs[i].len); // the varible to histify
    if(curO > histSlots-1){    // keep cur order from overrunning hist slots
      curO = histSlots-1;
    }
    branchHist[curO]++;
  }
}


//----------------------------------------------------------------------------
int overAny(){
  int over = -1;
  for(int i=0 ; i<numBranches ; i++){
    if(Bs[i].over() > 0){
      over = i;
    }
  }
  return over;
}


//----------------------------------------------------------------------------
// Delete branch from arbor, adjusting the ID#s and parent#s for the remaining
//  branches.  If a branch is being deleted, it must not have children
void killBranch(int id)
{
  //println("kill branch #" + id);
  Branch[] newBs = new Branch[maxBranches];
  for(int i=0 ; i<numBranches ; i++){
    if(Bs[i].parent > id){            // parents could have been created before the deleted branch
      Bs[i].parent = Bs[i].parent-1;
    }

    if(i > id){
      Bs[i].id = Bs[i].id-1;
      newBs[i-1] = Bs[i];
    } 
    else if(i < id){
      newBs[i] = Bs[i];
    }
  }

  Bs = newBs;
  numBranches--;
}


//----------------------------------------------------------------------------
// numVarSet: Sets the numerical value of simulation variables.
void numVarSet(float newNum){  
  if(adjMode == "pBranch"){
    pBranch = newNum;
    simVar = newNum;
  } 
  else if(adjMode == "pExt"){

  }
}


//----------------------------------------------------------------------------
// numVarInc: increments the numerical value of simulation variables.
void numVarInc(float inc){  
  if(adjMode == "pBranch"){
    float possible = 1.0 - pExt - pRet;
    if(pBranch+inc <= possible){
      pBranch = pBranch + inc;
    } 
    else {
      pBranch = possible;
    }
    simVar = pBranch;
  } 
  else if(adjMode == "pExt"){
    float possible = 1.0 - pBranch - pRet;
    if(pExt+inc <= possible){
      pExt = pExt + inc;
    } 
    else {
      pExt = possible;
    }
  }
  else if(adjMode == "pRet"){
    float possible = 1.0 - pBranch - pExt;
    if(pRet+inc <= possible){
      pRet = pRet + inc;
    } 
    else {
      pRet = possible;
    }
  }
}


//----------------------------------------------------------------------------
void initAttMap(int val){
  for(int i=0; i<attH; i++){
    for(int j=0; j<attW; j++){
      att[j][i] = val;
    }
  }
}


//----------------------------------------------------------------------------
void initInhibMap(int val){
  for(int i=0; i<inhibH; i++){
    for(int j=0; j<inhibW; j++){
      inhib[j][i] = val;
    }
  }
}






//----------------------------------------------------------------------------
//----------------------------------------------------------------------------
//----------------------------------------------------------------------------

/*  // bolds the growth probability currently under adjustment
 //String ptext = "| probabilities of... branching: " + pBranch + ", extension: " + pExt
 //  + ", retraction: " + pRet;
 //text(ptext, 50, 10);
 if(adjMode == "pBranch"){
 textFont(bold);
 text("branching: " + nf(pBranch, 1, 2), 10, 38);
 textFont(norm);
 text("extension: " + nf(pExt, 1, 2), 10, 52);
 text("retraction: " + nf(pRet, 1, 2), 10, 66);
 }
 if(adjMode == "pExt"){
 text("branching: " + nf(pBranch, 1, 2), 10, 38);
 textFont(bold);
 text("extension: " + nf(pExt, 1, 2), 10, 52);
 textFont(norm);
 text("retraction: " + nf(pRet, 1, 2), 10, 66);
 }
 if(adjMode == "pRet"){
 text("branching: " + nf(pBranch, 1, 2), 10, 38);
 text("extension: " + nf(pExt, 1, 2), 10, 52);
 textFont(bold);
 text("retraction: " + nf(pRet, 1, 2), 10, 66);
 textFont(norm);
 }
 *//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////