/**
 *
 * --Copyright notice-- 
 *
 * Copyright (c) Andy Evans. 
 * http://www.geog.leeds.ac.uk/people/a.evans/
 * This software is licensed under 'The Artistic License 2.0' which can be found at 
 * the Open Source Initiative website at... 
 * http://opensource.org/licenses/artistic-license-2.0
 * Please see disclaimers therein.
 *
 * The Standard Version source code, and associated documentation can be found at... 
 * Evans, A.J. (2007) 'Simple SIR model with Random Walk' MASS Group, University of Leeds, Britain 
 * [online] http://www.geog.leeds.ac.uk/courses/other/disease-modelling/ 
 *
 * --End of Copyright notice-- 
 *
**/

/**
 * Basic SIR disease model with random walk agents.
 * Runs as application or applet.<P>
 * Applet params as follows:<P>
 * "backgroundColor" example value: "#FFFFFF" - Background hex colour.<BR>
 * "textColor" example value: "#000000" - Foreground colour for text and box.<BR>
 * "initialSliderValues" example value: "100,100,0" - Sets the susceptibilitySlider, infectedTimeSlider, recoveryChanceSlider.<BR>
 * "lock" example value: "false" - whether to lock the values on the scrollbars.<BR>
 * @author Andy Evans: http://www.geog.leeds.ac.uk/people/a.evans/
 * @version 1.0 
 **/
import controlP5.*;
import java.awt.Toolkit; 
import java.awt.datatransfer.*;
import javax.swing.*;
import java.util.*;
import java.awt.Color;
import java.util.StringTokenizer;


// For sliders.
ControlP5 controlP5 = null;
String lock = "false";

// Fonts and GUI.
PFont titleFont;
PFont txtFont;
PFont boldFont;
int sliderSpacing = 90;
int txtColor = #C8DEF5;
int bgColor = #000000;
int sliderBackground = #003652;
int sliderForeground = #00698c;

String[] controlNames = {"susceptibilitySlider", "infectedTimeSlider", "recoveryChanceSlider"}; // These just used for slider names.
String[] controlLabels = {"Suscept\n-ibility(%)", "Infectious\nTime(s)", "Recovery\nChance(%)"};
int[] sliderColors = {#005a00, #640000, #1b1f9f, #333333}; // Note these are also the agent colours, matching the use of the scrollbars. Same order as agent types (below) so we can use types as index for colours.
int[] sliderDefaults = {50,50,50};  // Default starting values.

int[] colorBlindSliderColors = {#666666, #000000, #bbbbbb, #dddddd};
int colorBlindBgColor = #ffffff;
int colorBlindTxtColor = #000000;


// Agent array. Created with first agent and nulled when Q pressed.
// agents[0][] = x coordinate.
// agents[1][] = y coordinate.
// agents[2][] = type (below).
// agents[3][] = time until end of infection.
// agents[4][] = type at end of infection (INSUSCEPTIBLE or DEAD).
int agents[][] = null;              

// Agent types.
final int SUSCEPTIBLE = 0;
final int INFECTED = 1;
final int INSUSCEPTIBLE = 2;
final int DEAD = 3;

// Model running (i.e. Random walk motion).
boolean running = false;

// Dimensions of box for agents and agent size.
int stageLeft = 0;    
int stageTop = 0;
int stageWidth = 0;
int stageHeight = 0;
int circleSize = 10;
int lineStart = 33;

// For copy and paste.
Clipboard clipboard = null;
boolean[] keys = new boolean[526];
boolean copying = false;

// For recording data for copy and paste.
double timer = 0.0;
double oldTimer = 0.0;
String headerString = "Time,Total,Susceptible,Infected,Insusceptible,Dead\n";
String recordedData = headerString;





  /**
   * Draws the basic GUI and sets up clipboard if appropriate.
  **/
  void setup() {

    size(800, 400);
    
    // Initiate clipboard - fails for applets.
    try {
      clipboard = Toolkit.getDefaultToolkit().getSystemClipboard();
    } 
    catch (Exception e) {
      clipboard = null;
    }
  
    // Set up as applet or application
    if (clipboard != null)  {
      frame.setTitle("Disease modelling");
    } else {
       txtColor = getColor("textColor");
       bgColor = getColor("backgroundColor");
       sliderDefaults = getSliderDefaults("initialSliderValues");
       lock = getParameter("lock");
       if (lock == null) lock = "false";
       if (lock.equals("true")) {
         lineStart = 40;
       }
    }

    background(bgColor);
    

    // GUI
    smooth();
    txtFont = loadFont("MiriamFixed-12.vlw");
    boldFont = loadFont("CourierNewPS-BoldMT-12.vlw"); 


    // GUI elements
    controlP5 = new ControlP5(this);
    int xPosition = 0;


    // Sliders.
    for (int count = 0; count < controlNames.length; count++) {

      // Sliders   
      controlP5.addSlider(controlNames[count], 0, 100, sliderColors[count], xPosition + 40, 50-6, 14, height-120);
      Slider slider = (Slider)controlP5.controller(controlNames[count]);
      slider.setValue(sliderDefaults[count]);
      slider.setLabelVisible(false);
      if (lock.equals("true")) {
         slider.lock();  
      }
      
      slider.setColorActive(sliderForeground);
      slider.setColorBackground(sliderBackground);
      xPosition = xPosition + sliderSpacing;
    }

    // Now we know the size of the app, make the stage sizes.
    stageLeft = 40+90+90 + 100;
    stageTop = 50 - 6;
    stageWidth = (width - 20) - stageLeft;
    stageHeight = height - 120;
  
  }




  /**
  * This method returns a compressed colour int from a PARAM.<P>
  * The PARAM value should be a hash symbol and then a
  * standard web hex colour. If the PARAM is missing
  * the original instance variable defaults are returned.
  **/
  private int getColor(String param) {
    String colorString = getParameter(param);
    if (colorString == null) {
      if (param.equals("textColor"))
        return txtColor;
      else
        return bgColor;
      } else {
        try {
          int red = (Integer.decode("0x" + colorString.substring(1,3))).intValue();
          int green = (Integer.decode("0x" + colorString.substring(3,5))).intValue();
          int blue = (Integer.decode("0x" + colorString.substring(5))).intValue();
          return (new Color(red, green, blue)).getRGB();
        } catch (NumberFormatException nfe) {
          if (param.equals("textColor"))
            return txtColor;
          else
            return bgColor;
        }
      }

  } 
  
  
  
  
  /**
  * Returns an array of ints, comma separated in a param tag.
  **/
  private int[] getSliderDefaults(String param) {
    
    String defaultsString = getParameter(param);
    if (defaultsString == null) {
        return sliderDefaults;
      } else {
        
        try {
          
          StringTokenizer st = new StringTokenizer(defaultsString,",");
          int [] results = new int[st.countTokens()];
          for (int i = 0; i < results.length; i++) {
            results[i] = Integer.parseInt(st.nextToken());
            if (results[i] > 100) results[i] = 100;
            if (results[i] < 0) results[i] = 0;
          }
          return results;
          
        } catch (NumberFormatException nfe) {
          return sliderDefaults;
        }
      }
    
  }
 
  
  

 /**
  * Just used to reset slider colors to normal. 
  * These are set abnormally when altered using their own methods, injected below.
  **/
  public void mouseReleased() {

    for (int i = 0; i < controlNames.length; i++) {
      Slider s = (Slider)controlP5.controller(controlNames[i]);
      s.setColorBackground(sliderBackground);
      s.setColorActive(sliderForeground);
    }
    
  }




  /**
   * Changes color as slider moved.
  **/
  public void susceptibilitySlider(int theValue) {
    Slider s = (Slider)controlP5.controller("susceptibilitySlider");
    s.setColorActive(sliderColors[SUSCEPTIBLE]);
    s.setColorBackground(sliderColors[INSUSCEPTIBLE]);
  }




  /**
   * Changes color as slider moved.
  **/
  public void infectedTimeSlider(int theValue) {
    Slider s = (Slider)controlP5.controller("infectedTimeSlider");
    s.setColorActive(sliderColors[INFECTED]);
    s.setColorBackground(sliderBackground);
  }




  /**
   * Changes color as slider moved.
  **/
  public void recoveryChanceSlider(int theValue) {
    Slider s = (Slider)controlP5.controller("recoveryChanceSlider");
    s.setColorActive(sliderColors[INSUSCEPTIBLE]);
    s.setColorBackground(sliderColors[DEAD]);
  }




  /**
  * Adds agents, with types depending on order and scrollbars.
  * Also notes any changes and calls routine to update potentially copied data if 
  * there have beeen any.
  **/
  void mousePressed() {
  
    if(((mouseX < stageLeft) || (mouseX > stageLeft + stageWidth)) ||
      ((mouseY < stageTop) || (mouseY > stageTop + stageHeight))) return; 

    if ((agents == null) || (mouseButton == RIGHT)) {
      
      addAgent(mouseX, mouseY, INFECTED); 
    
    } else {
    
      Slider slider = (Slider)controlP5.controller("susceptibilitySlider");
      int value = int(slider.value());
    
      if (Math.random()*100 < value) {
        addAgent(mouseX, mouseY, SUSCEPTIBLE); 
      } else {
        addAgent(mouseX, mouseY, INSUSCEPTIBLE); 
      }
  
    }  
    
    recordChanges();
    
  }




  /**
  * Adds agents, with types depending on order and scrollbars.
  **/
  private void addAgent(int x, int y, int type) {
  
    // If first agent, make them infected. Mainly for those without a right mouse button.
    if (agents == null) {
      
      agents = new int[5][1];
      agents[0][0] = x;
      agents[1][0] = y;
      agents[2][0] = type;
      
      // Set potential infected time.
      Slider slider = (Slider)controlP5.controller("infectedTimeSlider");
      int value = int(slider.value());
      agents[3][0] = value * (int)frameRate;
      
      // Set potential result of infection.
      // Note we sort this out now so agent keeps probability at time of formation.  
      Slider slider2 = (Slider)controlP5.controller("recoveryChanceSlider");
      int value2 = int(slider2.value());
      int rand = (int)(Math.random()*100);
      agents[4][0] = (rand > value2)? DEAD:INSUSCEPTIBLE; 
      
      timer = 0.0; // Zero time incase they've let the model run with no agents since start / Q pushed.
     
      return;
    }   
  
    // Otherwise increase the size of the agents array and add another.
    int[][] temparray = new int [5][agents[0].length];
    
    for (int i = 0; i < agents[0].length; i++) {
      temparray[0][i] = agents[0][i];
      temparray[1][i] = agents[1][i];
      temparray[2][i] = agents[2][i];
      temparray[3][i] = agents[3][i];
      temparray[4][i] = agents[4][i];
    }
    
    agents = new int[5][agents[0].length + 1];
    
    for (int i = 0; i < temparray[0].length; i++) {
      agents[0][i] = temparray[0][i];
      agents[1][i] = temparray[1][i];
      agents[2][i] = temparray[2][i];
      agents[3][i] = temparray[3][i];
      agents[4][i] = temparray[4][i];
    }
  
    // New agent.
    agents[0][agents[0].length - 1] = x;
    agents[1][agents[1].length - 1] = y;
    agents[2][agents[2].length - 1] = type;
    
    // Set potential infected time.
    Slider slider = (Slider)controlP5.controller("infectedTimeSlider");
    int value = int(slider.value());
    agents[3][agents[2].length - 1] = value * (int)frameRate;
    
    // Set potential result of infection.
    // Note we sort this out now so agent keeps probability at time of formation.
    Slider slider2 = (Slider)controlP5.controller("recoveryChanceSlider");
    int value2 = int(slider2.value());
    int rand = (int)(Math.random()*100);
    agents[4][agents[2].length - 1] = (rand > value2)? DEAD:INSUSCEPTIBLE;    
    
    return;
  }





  /**
   * Random walk motion. Very basic.
  **/
  private void moveRandomly() {
 
    double xRand = 0;
    double yRand = 0;
  
    for (int i = 0; i < agents[0].length; i++) {
    
      if (agents[2][i] == DEAD) continue; // Skip the dead.
    
      xRand = Math.random();
      yRand = Math.random();
    
      if (xRand > 0.5) agents[0][i]++;
      else agents[0][i]--;
    
      if (yRand > 0.5) agents[1][i]++;
      else agents[1][i]--;

      // Make sure they don't wander off stage.
      if (agents[0][i] < stageLeft) agents[0][i]++;
      if (agents[0][i] > stageLeft + stageWidth) agents[0][i]--;
      if (agents[1][i] < stageTop) agents[1][i]++;
      if (agents[1][i] > stageTop + stageHeight) agents[1][i]--;
    
    }
  
  }




  /**
  * Check for collisions. Again, very basic version.
  **/
  private void checkCollisions() {
  
    int xDiff = 0;
    int yDiff = 0;
  
    boolean statesChanged = false;
  
    for (int i = 0; i < agents[0].length; i++) {
  
      if (agents[2][i] == DEAD) continue; // Skip the dead
    
      for (int j = i; j < agents[0].length; j++) { // NB Only test against those we haven't yet tested as agent[i]
       
        xDiff = Math.abs(agents[0][i] - agents[0][j]);
        yDiff = Math.abs(agents[1][i] - agents[1][j]);
        
        if ((xDiff < circleSize) && (yDiff < circleSize)) {
        
          if ((agents[2][i] == INFECTED) && (agents[2][j] == SUSCEPTIBLE)) {
            agents[2][j] = INFECTED;
            statesChanged = true;
          }
          if ((agents[2][j] == INFECTED)&& (agents[2][i] == SUSCEPTIBLE)) {
            agents[2][i] = INFECTED;
            statesChanged = true;
          }
        
        }
      
      }
      
      // Reduce time recovered (agents[3]) and kill or save  as stored in agents[4].
      if (agents[2][i] == INFECTED) {
        agents[3][i]--; 
        if (agents[3][i] < 1) { 
          agents[2][i] = agents[4][i];
          statesChanged = true;
        }
      }
   
    }
 
    // If any agents have changed, add new summary of data to data to be copied and pasted.
    if (statesChanged) { 
      recordChanges();
      statesChanged = false;
    }    
  
  }




  /**
   * Set us in copy mode and copies to clipboard if CTRL-C pressed.
  **/
  void keyPressed() {
  
    keys[keyCode] = true;
    if (checkKey(CONTROL) && checkKey(KeyEvent.VK_C)) {
      // Enbolden the copy help message (see draw()).
      copying = true;

      // Copy the recordedData.
      StringSelection s = new StringSelection(recordedData);
      if (clipboard != null) clipboard.setContents(s, s);
      return;
    }

    // Run or stop
    if (checkKey(KeyEvent.VK_SPACE)) {
      running = !running;
      return;
    }

    // Clear everything
    if (checkKey(KeyEvent.VK_Q)) {
      agents = null;
      timer = 0.0;
      recordedData = headerString; // First agent added in mousePressed.
      return;
    }


    // Colorblind option
    if (checkKey(KeyEvent.VK_G)) {
      
      int tempBgColor = bgColor;
      bgColor = colorBlindBgColor;
      colorBlindBgColor = tempBgColor;
      
      int tempTxtColor = txtColor;
      txtColor = colorBlindTxtColor;
      colorBlindTxtColor = tempTxtColor;
      
      int[] tempStoredColors = new int[sliderColors.length];
      System.arraycopy(sliderColors,0,tempStoredColors,0,sliderColors.length);
      System.arraycopy(colorBlindSliderColors,0,sliderColors,0,sliderColors.length);
      System.arraycopy(tempStoredColors,0,colorBlindSliderColors,0,sliderColors.length);
      
      Slider s = (Slider)controlP5.controller("susceptibilitySlider");
      s.setColorActive(sliderForeground);
      s.setColorBackground(sliderBackground);
      s.updateEvents();
      s = (Slider)controlP5.controller("infectedTimeSlider");
      s.setColorActive(sliderForeground);
      s.setColorBackground(sliderBackground);
      s = (Slider)controlP5.controller("recoveryChanceSlider");
      s.setColorActive(sliderForeground);
      s.setColorBackground(sliderBackground);
      background(bgColor);

      return;
    }


  }




  /**
   * Set us out of copy mode.
  **/
  void keyReleased() { 
    
    keys[keyCode] = false; 
  
    // Unenbolden the copy help message (see draw()).
    copying = false;
  
  }




  /**
   * CTRL checking code from http://wiki.processing.org/w/Multiple_key_presses
  **/
  boolean checkKey(int k) {
    if (keys.length >= k) {
      return keys[k];
    }
    return false;
  }




  /**
  * Adds to store of data ready for copying and pasting.
  **/
  private void recordChanges() {
   
    // recordedData reinitialised whenever cleared.
    // Timer reinitialised in same place.
    // recordChanges called when agents added or collide, or recover/die.
    
    int [] typeCounts = new int[4]; 
    
    // Count agents of different types.
    for (int i = 0; i < agents[0].length; i++) {
      typeCounts[agents[2][i]]++;
    }
    
    // Construct other variables.
    double time = (timer/frameRate);
    int total = typeCounts[SUSCEPTIBLE] + typeCounts[INFECTED] + typeCounts[INSUSCEPTIBLE] + typeCounts[DEAD];
    
    // If the user hasn't run the model between changes, usually because they are adding lots of agents, and there is data.
    // Remove the last entry so they don't get a record for each agent added (NB they do get this if the model is running, 
    // but then they're adding them over time, here they're not, the time would be the same for each agent).
    if ((timer == oldTimer) && (recordedData.length() > headerString.length())) { 
      recordedData = recordedData.substring(0,recordedData.lastIndexOf("\n", recordedData.length() - 2) + 1);  
    }
    
    // Add to the stored data string.
    recordedData = recordedData + 
      time + "," + 
      total + "," + 
      typeCounts[SUSCEPTIBLE] + "," +  
      typeCounts[INFECTED] + "," +
      typeCounts[INSUSCEPTIBLE] + "," +  
      typeCounts[DEAD] + "\n"; 

    // Reset the timer to check whether they've run the model.
    oldTimer = timer;
    
  }




  /**
   * Draws the GUI with new color.
  **/
  void draw() {
  
    // Increment time running since first run started / clear key pushed.
    if (running) timer++;
  
  
  
  
   // GUI
   
   
    textFont(txtFont);
    textSize(12);
    smooth();
    
    
    // GUI elements
    int xPosition = 0;



    // Sliders' text down side.
    textAlign(RIGHT);
    
    for (int count = 0; count < controlNames.length; count++) {

      fill(bgColor);
      noStroke();
      rect(xPosition + 15, 42, 24, height-110);
      fill(txtColor);
      
      // Decimal numbers
      for (int i = 100; i >= 0; i -= 10) {
        text(i, xPosition + 39, 50 + (100-i)*(((height-120.0)/100.0)));
      }   

      xPosition = xPosition + sliderSpacing;
    }


    // Little flashes to show column colours.
    noStroke();

    fill(sliderColors[0]);
    rect(40, (50-6)+(height-120)+5, 14, 1);

    fill(sliderColors[1]);
    rect(40+90, (50-6)+(height-120)+5, 14, 1);

    fill(sliderColors[2]);
    rect(40+90+90, (50-6)+(height-120)+5, 14, 1);

  
  
  
  
  
  
    // Paint over the old values from the sliders.
    fill(bgColor);
    noStroke();
    rect(0, height-70, width, 60);
  
    // Paint the new values from the sliders.
    fill(txtColor);
    xPosition = 0;
    Slider slider = null;

    int value = 0;
 
    // Loop through sliders.
    for (int i = 0; i < controlNames.length; i++) { 

      // Get slider value.
      slider = (Slider)controlP5.controller(controlNames[i]);
      value = int(slider.value());

      // Print value on GUI.
      textAlign(LEFT);
      text(value, xPosition + 40, 50 + (height-100.0));
      text(controlLabels[i], xPosition + 40, 50 + (height-80.0));
      xPosition = xPosition + 90;
    
    }


    // Paint the old stage over.

    fill(bgColor);
    rect(stageLeft + 1, stageTop + 1, stageWidth - 1, stageHeight - 1);
   
  
    // Paint the agents.
    if (agents != null) {  
       for (int i = 0; i < agents[0].length; i++) {
         fill(sliderColors[agents[2][i]]);
         ellipse(agents[0][i],agents[1][i],circleSize,circleSize);
       }
    }

  
    // Paint the stage box, and some black boxes around it to stop the agents appear 
    // half outside the boundary when their centres are up against it.
    stroke(txtColor);
    line(stageLeft, stageTop, stageLeft, stageTop + stageHeight);
    line(stageLeft, stageTop, stageLeft + stageWidth, stageTop);
    line(stageLeft + stageWidth, stageTop, stageLeft + stageWidth, stageTop + stageHeight);
    line(stageLeft, stageTop + stageHeight, stageLeft + stageWidth, stageTop + stageHeight);
    fill(bgColor);
    stroke(bgColor);
    rect(stageLeft - circleSize, stageTop - circleSize, circleSize - 1, stageHeight + circleSize + circleSize); // left
    rect(stageLeft - circleSize, stageTop - circleSize, stageWidth + circleSize + circleSize, circleSize - 1); // top
    rect(stageLeft + stageWidth + 1, stageTop - circleSize, circleSize, stageHeight + circleSize + circleSize); // right
    rect(stageLeft - circleSize, stageTop + stageHeight + 1, stageWidth + circleSize + circleSize, circleSize); // bottom
    fill(txtColor);
  
  
    // Move the agents read for next step.
    if ((agents != null)&&(running)) {
      moveRandomly();
      checkCollisions();
    }
  

    
  
    // Add instructions.
    
    text("Click to add people; 1st person infected, also right mouse button.", stageLeft, lineStart + int((height-90.0)));
    
    if (lock.equals("true")) {
      text("Spacebar runs/stops. Q clears. G for greyscale.", stageLeft, lineStart + 15 + int((height-90.0)));
    } else {
      text("Scrollbars change new people only. Spacebar runs/stops. Q clears. ", stageLeft, lineStart + 15 + int((height-90.0)));
      if (clipboard != null) {
        text("G for greyscale.", stageLeft, lineStart + 45 + int((height-90.0)));
      } else {
        text("G for greyscale.", stageLeft, lineStart + 30 + int((height-90.0)));
      }
    }
    // If we're in copy mode, bold up the next text to give visual clue.
    if (clipboard != null) {
      if (copying) {
        textFont(boldFont);
      } else {
        textFont(txtFont);
      }
      text("Ctrl+C copies infection data. Use Ctrl+V in other apps to paste.", stageLeft, lineStart + 30 + int((height-90.0)));
      
      textFont(txtFont);
    }

}