<iframe width='500px' height='400px' src='https://nest.ktbyte.com/nest#292271' allowfullscreen></iframe>

int SUSCEPTIBLE = 0, 
  INFECTIOUS = 1, 
  REMOVED = 2;
int NUMBER_PEOPLE = 100;
int NUMBER_INFECTED = 2;
int ASYMPTOMATIC_PROBABILITY = 20;
// int PERCENT_ADHERING_SOCIAL_DISTANCE = 70;
int PERCENT_ADHERING_SOCIAL_DISTANCE = 100;
boolean USE_SOCIAL_DISTANCING = false;
int HIGHEST_PROBABILIY_POP_B = 100;
// int HIGHEST_PROBABILIY_POP_B = 0;
int DEATH_PROBABILITY_POP_A = 0;
int DEATH_PROBABILITY_POP_B = 20;
SIRModel model = new SIRModel(500, 350, NUMBER_PEOPLE, NUMBER_INFECTED);
SIRModel visibleModel = new SIRModel(500, 350, model);
int INFECT_DISTANCE = 8;
int INFECT_PROBABILITY = 10;

boolean TESTING_QUARANTINE_APPLIES_TO_ASYMTOMATIC = false;
float WALK_SPEED = 1;
float FAST_WALK_SPEED = 5;
int TESTING_RADIUS = 50;

int AUTO_QUARANTINE_SYMPTOMATIC_AFTER_FRAMES = 60*50;

//Each 60 frames, a second, is 4 days let's say
int RECOVERY_FRAME_COUNT = 30 * 14; //14 days to recovery
int FRAMES_TO_SHOW_INFECTION = 60 * 2;
color COLOR_S = color(100, 150, 255), 
  COLOR_I = color(255, 150, 100), 
  COLOR_R = color(150, 255, 100), 
  COLOR_D = color(0);
color COLOR_ASYMTOMATIC = color(255, 0, 0);

int framesPerChartColumn = 10;
int FRAMES_PER_CHART_TICK = 30 * 7;
int DAYS_PER_TICK = 7;

void setup() {
  size(1400, 900);
}
void draw() {
  background(255);
  model.display(50, 50);
  model.step();
  model.test(50, 50);
  model.computeChart();
  model.chart(600, 50, 600, 350);
  visibleModel.testX = model.testX;
  visibleModel.testY = model.testY;
  visibleModel.display(50, 500);
  visibleModel.time++;
  visibleModel.computeChart();
  visibleModel.chart(600, 500, 600, 350);
}
class Person {
  int state = SUSCEPTIBLE, 
    w, 
    h;
  boolean dead = false;
  int infectedFrames = 0;
  boolean popB;

  boolean quarantine = false; // When quarantined move to the top left corner
  boolean asymptomatic = random(0, 100) < ASYMPTOMATIC_PROBABILITY;
  boolean adheres_to_social_distancing = random(0, 100) < PERCENT_ADHERING_SOCIAL_DISTANCE;
  int transmissions = 0;
  Set<Person> possibleTransmissions = new HashSet<Person>();
  Person(int w, int h) {
    this.w = w;
    this.h = h;
    x = random(0, w);
    y = random(0, h);
    if (random(0, 100) < x*1.0/w*HIGHEST_PROBABILIY_POP_B) {
      popB = true;
    }
  }
  float x, 
    y, 
    angle = random(0, 2 * PI), 
    speed = WALK_SPEED, 
    rotatePerStep = PI / 10;
  int showInfectFrames = 0;
  void infect(Person source) {
    state = INFECTIOUS;
    showInfectFrames = FRAMES_TO_SHOW_INFECTION;
    source.transmissions++;
  }
  void addPotentialInfection(Person other) {
    possibleTransmissions.add(other);
  }
  void display() {
    if (dead) {
      fill(0);
    } else if (state == 0) fill(COLOR_S);
    else if (state == 1) {
      fill(COLOR_I);
      if (asymptomatic) fill(COLOR_ASYMTOMATIC);
    } else if (state == 2) fill(COLOR_R);
    noStroke();
    if (!dead) {
      rectMode(CENTER);
      if (popB) {
        stroke(0);
        rect(x, y, 8, 8);
      } else {
        ellipse(x, y, 8, 8);
      }
      noStroke();
    } else {
      rectMode(CENTER);
      stroke(0);
      line(x-2, y-2, x+2, y+2);
      line(x-2, y+2, x+2, y-2);
      //rect(x, y, 5, 5);
    }
    if (showInfectFrames > 0) {
      noFill();
      stroke(255, 0, 0, map(showInfectFrames, 0, FRAMES_TO_SHOW_INFECTION, 0, 255));
      ellipse(x, y, INFECT_DISTANCE * 2, INFECT_DISTANCE * 2);
      showInfectFrames--;
    }
  }
  void step(int testX, int testY, ArrayList<Person> people) {
    if (dist(x, y, testX, testY) < TESTING_RADIUS && !dead) {
      if (!visibleModel.people.contains(this)) {
        visibleModel.people.add(this);
      }
    }
    if (state == INFECTIOUS && infectedFrames > AUTO_QUARANTINE_SYMPTOMATIC_AFTER_FRAMES && !asymptomatic) {
      quarantine = true;
      goToQuarantine();
    } else if (state == INFECTIOUS && (quarantine || (dist(x, y, testX, testY) < TESTING_RADIUS && (!asymptomatic || TESTING_QUARANTINE_APPLIES_TO_ASYMTOMATIC)))) {
      quarantine = true;
      goToQuarantine();
    } else if (!dead) {
      if (USE_SOCIAL_DISTANCING && adheres_to_social_distancing) avoidClosest(people);
      else randomWalk();
    }
    //else if (dead && dist(x, y, testX, testY) < TESTING_RADIUS) {
    //  if (!visibleModel.people.contains(this)) {
    //    visibleModel.people.add(this);
    //  }
    //}
    if (state == INFECTIOUS) infectedFrames++;
    if (infectedFrames > RECOVERY_FRAME_COUNT && state != REMOVED) {
      state = REMOVED;
      //some probability of dying
      if (!popB) {
        if (random(0, 100) < DEATH_PROBABILITY_POP_A) {
          dead = true;
        }
      } else {
        if (random(0, 100) < DEATH_PROBABILITY_POP_B) {
          dead = true;
        }
      }
    }
  }
  float quarantineX = random(-50, -20), quarantineY = random(-50, -20);
  void goToQuarantine() {
    angle = atan2(quarantineY - y, quarantineX - x);
    if (dist(x, y, quarantineX, quarantineY) < 50) speed = WALK_SPEED;
    else speed = FAST_WALK_SPEED;
    x += speed * cos(angle);
    y += speed * sin(angle);
  }
  void randomWalk() {
    speed = WALK_SPEED;
    x += speed * cos(angle);
    y += speed * sin(angle);
    angle += random( - rotatePerStep, rotatePerStep);
    if (x > w || x < 0 || y < 0 || y > h) angle = random( - PI, PI);
    x = constrain(x, 0, w);
    y = constrain(y, 0, h);
    //if at edges, turn faster
  }
  void avoidClosest(ArrayList<Person> people) {
    Person closest = null;
    for (Person p : people) {
      if (p != this) {
        if (closest == null) closest = p;
        else if (dist(p.x, p.y, x, y) < dist(closest.x, closest.y, x, y)) closest = p;
      }
    }
    if (closest != null) {
      angle = atan2(y-closest.y, x-closest.x);
      randomWalk();
    }
  }
}
class SIRHistory {
  int s;
  int i, 
    a;
  int r, 
    d;
}
class SIRModel {
  int time = 0;
  ArrayList<SIRHistory> history = new ArrayList<SIRHistory>();
  SIRModel originalModel = null;
  int peakInfection = 0, peakTotal = 1;
  int testX = -TESTING_RADIUS*10, testY = -TESTING_RADIUS*10;
  SIRModel(int w, int h, SIRModel originalModel) {
    this.w = w;
    this.h = h;
    this.originalModel = originalModel;
  }
  SIRModel(int w, int h, int n, int infected) {
    this.w = w;
    this.h = h;
    for (int i = 0; i < n; i++) {
      Person p = new Person(w, h);
      people.add(p);
      if (i < infected) p.state = INFECTIOUS;
    }
  }
  int w, 
    h;
  ArrayList < Person > people = new ArrayList < Person > ();

  //top left x y coordinate
  void display(int xOffset, int yOffset) {
    pushMatrix();
    translate(xOffset, yOffset);
    rectMode(CORNER);
    fill(255);
    stroke(0);
    rect(0, 0, w, h);
    for (Person p : people) {
      p.display();
    }

    noStroke();
    if (testX > 0 && testY > 0) {
      fill(0, 50);
      ellipse(testX, testY, TESTING_RADIUS*2, TESTING_RADIUS*2);
    }
    float quarantineX = random(-50, -20), quarantineY = random(-50, -20);
    fill(0,50);
    ellipse((-50+20)/2,(-50+20)/2,(-20+50),(-20+50));
    popMatrix();
  }
  void test(int xOffset, int yOffset) {
    if (mouseX - xOffset > 0 && mouseX - xOffset < w && mouseY - yOffset > 0 && mouseY - yOffset < h) {
      testX = mouseX - xOffset;
      testY = mouseY - yOffset;
    } else testX=testY=-TESTING_RADIUS*10;
  }
  void chart(int xOffset, int yOffset, int chartW, int chartH) {

    pushMatrix();
    translate(xOffset, yOffset);
    rectMode(CORNER);
    fill(255);
    stroke(0);
    rect(0, 0, chartW, chartH);

    if (history.size() == 0) {
      popMatrix();
      return;
    }
    float barWidth = 1.0 * chartW / history.size();

    int d = 0, 
      t = 0, 
      cfrifrPercent = 0;
      
    for (int i = 0; i < history.size(); i++) {
      boolean lastColumn = i == history.size() - 1;
      SIRHistory sir = history.get(i);
      int total = sir.s + sir.i + sir.r;
      if(total == 0) continue;
      float x = barWidth * i - 1;
      float heightS = sir.s * 1.0 / total * chartH;
      float heightRalive = (sir.r - sir.d) * 1.0 / total * chartH;
      float heightRdead = sir.d * 1.0 / total * chartH;
      textAlign(LEFT, CENTER);
      noStroke();
      fill(COLOR_D);
      rect(x, 0, barWidth + 1, heightRdead);
      if (lastColumn && heightRdead > 0) text("DEAD", chartW + 3, heightRdead / 2.0);
      fill(COLOR_R);
      rect(x, heightRdead, barWidth + 1, heightRalive);
      if (lastColumn) text("REMOVED", chartW + 3, heightRdead + heightRalive / 2.0);

      float heightR = heightRalive + heightRdead;
      fill(COLOR_S);
      rect(x, heightR, barWidth + 1, heightS);
      if (lastColumn) text("SUSCEPTIBLE", chartW + 3, heightR + heightS / 2.0);

      float heightIasymptomatic = (sir.a) * 1.0 / total * chartH;
      float heightIsymptomatic = (sir.i - sir.a) * 1.0 / total * chartH;
      fill(COLOR_I);
      rect(x, heightR + heightS, barWidth + 1, heightIsymptomatic);
      if (lastColumn) {
          text("INFECTIOUS-SYMPTOMATIC", chartW + 3, heightR + heightS + heightIsymptomatic / 2.0);
          if(sir.i > peakInfection) {
              peakInfection = sir.i;
              peakTotal = total;
          }
      }

      fill(COLOR_ASYMTOMATIC);
      fill(255, 0, 0);
      rect(x, heightR + heightS + heightIsymptomatic, barWidth + 1, heightIasymptomatic);
      if (lastColumn && sir.a > 0) text("ASYMPTOMATIC", chartW + 3, heightR + heightS + heightIsymptomatic + heightIasymptomatic / 2.0);

      if ((i * framesPerChartColumn) % (FRAMES_PER_CHART_TICK) == 0) {
        stroke(0);
        line(x, chartH, x, chartH + 10);
        fill(0);
        String txt = "" + (i * framesPerChartColumn) / FRAMES_PER_CHART_TICK * DAYS_PER_TICK;
        textAlign(CENTER, CENTER);
        text(txt, x, chartH + 25);
      }
      if (lastColumn) {
        d = sir.d;
        t = total;
        cfrifrPercent = (int)(sir.d * 1.0 / (sir.s + sir.i + sir.r) * 100);
      }
    }

    fill(0);
    
    float sumR = 0, sumR0 = 0, countR = 0;
    for(Person p : people) {
        if(p.transmissions > 0) {
            sumR += p.transmissions;
            sumR0 += p.possibleTransmissions.size();
            countR++;
        }
    }
    
    String stats = "CFR/IFR: " + d + "/" + t + "=" + cfrifrPercent + "%,";
    stats += " Peak: "+(int)(peakInfection * 100.0 / peakTotal)+"%,";
    if(people.size()==NUMBER_PEOPLE && countR > 0) {
        stats += " R: "+((int)(sumR * 10.0/ countR))/10.0+",";
        String str2 = ""+((int)(sumR0 * 10.0/ countR))/10.0;
        if(!str2.contains(".")) str2 += ".0";
        stats += " R0: "+str2+",";
    }
    stats += " Days: "+(history.size()* framesPerChartColumn * DAYS_PER_TICK / FRAMES_PER_CHART_TICK);
    
    textAlign(CENTER,CENTER);
    text(stats, chartW / 2, -20);

    popMatrix();
  }
  void step() {
    time++;
    for (Person p : people) {
      p.step(testX, testY, people);
      if (p.speed == FAST_WALK_SPEED) continue;
      //Check for infection
      for (Person other : people) {
        if (p != other && random(0, 100) < INFECT_PROBABILITY) {
          if (dist(p.x, p.y, other.x, other.y) < INFECT_DISTANCE) {
            if (p.state == INFECTIOUS && other.state == SUSCEPTIBLE) {
              other.infect(p);
            }
            if(p.state == INFECTIOUS) p.addPotentialInfection(other);
          }
        }
      }
    }
  }
  void computeChart() {
    // time++;
    if (time % framesPerChartColumn == 0) {
      SIRHistory now = new SIRHistory();
      for (Person p : people) {
        if (p.state == SUSCEPTIBLE) now.s++;
        if (p.state == INFECTIOUS) {
          if (p.asymptomatic) now.a++;
          now.i++;
        }
        if (p.state == REMOVED) now.r++;
        if (p.dead) now.d++;
      }

      //if (now.i > QUARANTINE_AFTER_INFECTED_NUMBER) quarantineEnabled = true;

      if (now.i > 0 || (originalModel != null && originalModel.history.size() > history.size())) history.add(now);
      else {
        //Check if this is the last one
        if (history.get(history.size() - 1).i > 0) history.add(now);
      }
    }
  }
}
void mousePressed() {
  USE_SOCIAL_DISTANCING = !USE_SOCIAL_DISTANCING;
}