boolean enableBestN = true;
boolean enableImprovReq = true;
boolean doOutputData = false;
String[] toOutput;

Particle[] particles;
int particleCount = 100;
int dimX = 600;
int dimY = 600;
int bottomPanel = 100;

int maxAge = 1000;
int age = 0;

float tx, ty; // The target

int nBest = 5;
PriorityList bestIndices;
float bestWeighting = 0.5;
//int[] bestIndices;

/*int bestParticleIndex;
float bestFitness = 100000000;*/

//float nalpha1  = 0.075;
//float nalpha12 = 0.150;
float nalpha1 = 0.015;
float nalpha12 = 0.030;
float nalpha2  = 0.010;
float nalpha22 = 0.020;
float nmaxVelocity  = 0.5;
float nmaxVelocity2 = 1;

float reqdImprv = .200; // Required improvement
float reqdEpochs = 10; // Number of epochs in which improvement is required
int reqdProb = 20; // Probation period for particles.

boolean introducing = true;

boolean locked = false; // For buttons
MyButton incMaxVelocity, decMaxVelocity;
MyButton incAlpha1, decAlpha1;
MyButton incAlpha2, decAlpha2;
MyButton incNumBest, decNumBest;
color buttonColour = color(128);
color buttonHighlight = color(160);

void setup() {
  if (!enableBestN) {
    nBest = 1;
    bestWeighting = 1;
  }
  if (!enableImprovReq) {
    reqdEpochs = maxAge + 1;
  }
  
  size(600, 700);
  smooth();
  background(0);
  frameRate(30);
  
  particles = new Particle[particleCount];
  
  if (doOutputData) {
    toOutput = new String[maxAge + 1];
  }
  
  textFont(loadFont("Calibri-12.vlw"));
   
  decMaxVelocity = new MyButton(180, dimY + 4, 12, 12, buttonColour, buttonHighlight, " -");
  incMaxVelocity = new MyButton(260, dimY + 4, 12, 12, buttonColour, buttonHighlight, " +");
  decAlpha1 = new MyButton(180, dimY + 28, 12, 12, buttonColour, buttonHighlight, " -");
  incAlpha1 = new MyButton(260, dimY + 28, 12, 12, buttonColour, buttonHighlight, " +"); 
  decAlpha2 = new MyButton(180, dimY + 52, 12, 12, buttonColour, buttonHighlight, " -");
  incAlpha2 = new MyButton(260, dimY + 52, 12, 12, buttonColour, buttonHighlight, " +"); 
  if (enableBestN) {
    decNumBest = new MyButton(180, dimY + 76, 12, 12, buttonColour, buttonHighlight, " -");
    incNumBest = new MyButton(260, dimY + 76, 12, 12, buttonColour, buttonHighlight, " +");
  }
  drawButtons();
}

void reset() {
  background(0);
  for (int i = 0; i < particleCount; i++) {
    particles[i] = new Particle(i, -1, -1);
  }
  age = 0;
  drawButtons();
  /*bestParticleIndex = 0;
  bestFitness = 100000000;*/
  loop();
}

void mousePressed() {
  if (mouseY > dimY) {
    if (incMaxVelocity.pressed()) {
      nmaxVelocity += 0.5;
      nmaxVelocity2 = nmaxVelocity * 2;
    } else if (decMaxVelocity.pressed()) {
      nmaxVelocity -= 0.5;
      nmaxVelocity = max(1, nmaxVelocity);
      nmaxVelocity2 = nmaxVelocity * 2;
    } else if (incAlpha1.pressed()) {
      nalpha1 += 0.005;
      nalpha12 = nalpha1 * 2;
    } else if (decAlpha1.pressed()) {
      nalpha1 -= 0.005;
      nalpha1 = max(0.005, nalpha1);
      nalpha12 = nalpha1 * 2;
    } else if (incAlpha2.pressed()) {
      nalpha2 += 0.005;
      nalpha22 = nalpha2 * 2;
    } else if (decAlpha2.pressed()) {
      nalpha2 -= 0.005;
      nalpha2 = max(0.005, nalpha2);
      nalpha22 = nalpha2 * 2;
    } else if (enableBestN) {
      if (incNumBest.pressed()) {
        nBest += 1;
      } else if (decNumBest.pressed()) {
        nBest -= 1;
        nBest = max(1, nBest);
      }
    }
  } else {
    tx = mouseX;
    ty = mouseY;
    introducing = false;
    reset();
  }
}

void drawButtons() {
  fill(80);
  noStroke();
  rect(0, dimY, dimX, bottomPanel);
  if (!locked) {
    incMaxVelocity.update();
    decMaxVelocity.update();
    incAlpha1.update();
    decAlpha1.update();
    incAlpha2.update();
    decAlpha2.update();
    if (enableBestN) {
      incNumBest.update();
      decNumBest.update();
    }
  } else {
    locked = false;
  }
  fill(255);
  text("Maximum Velocity:", 12, dimY + 12);
  text("Alpha 1:", 12, dimY + 36);
  text("Alpha 2:", 12, dimY + 60);
  if (enableBestN) {
    text("# of Best:", 12, dimY + 84);
  }
  fill(255, 204, 0);
  text(nf(nmaxVelocity, 1, 0), 216, dimY + 12);
  text(nf(nalpha1, 1, 3), 216, dimY + 36);
  text(nf(nalpha2, 1, 3), 216, dimY + 60);
  if (enableBestN) {
    text(nf(nBest, 1, 0), 216, dimY + 84);
  }
}

void draw() {
  // Control panel
  fill(80);
  noStroke();
  rect(0, dimY, dimX, bottomPanel);
  
  if (introducing) {
    int rowHeight = 24;
    int y = 24;
    int x = 12;
    int i = 0;
    background(0);
    fill(255, 204, 0);
    text("Particle Swarm Optimisation (Enhanced)", x, y + rowHeight * i++);
    fill(255, 192);
    i++;
    text("This is an enhanced PSO. General points to note:", x, y + rowHeight * i++);
    text("- Particles are assigned a random maximum velocity between " + nmaxVelocity + " and " + nmaxVelocity2 + ".", x, y + rowHeight * i++);
    text("- Particles move towards the global best with a random alpha between " + nalpha1 + " and " + nalpha12 + ".", x, y + rowHeight * i++);
    text("- Particles move towards the local best with a random alpha between " + nalpha2 + " and " + nalpha22 + ".", x, y + rowHeight * i++);
    text("- Particle fitness is defined as sqrt((tx - x)^2 + (ty - y)^2).", x, y + rowHeight * i++);
    text("- The simulation runs for " + maxAge + " epochs.", x, y + rowHeight * i++);
    text("- You may restart the simulation at any time by clicking; the target will be at the location of your click.", x, y + rowHeight * i++);
    i++;
    fill(255);
    text("Enhancements:", x, y + rowHeight * i++);
    fill(255, 192);
    if (enableBestN) {
      text("- There are now " + nBest + " 'best' particles which attract all others.", x, y + rowHeight * i++);
      text("- The 'best' particles are highlighted in red.", x, y + rowHeight * i++);
      text("- Particles which are marked as 'best' are not influenced by themselves.", x, y + rowHeight * i++);
      text("- Particles which are marked as 'best' have influence from other 'best' particles weighted by " + bestWeighting + ",", x, y + rowHeight * i++);
      text("  and have their own local influence weighted by " + bestWeighting + "^-1.", x, y + rowHeight * i++);
    }
    if (enableImprovReq) {
      text("- If a particle's fitness does not improve by a fraction of " + reqdImprv + " in " + reqdEpochs + "epochs,", x, y + rowHeight * i++);
      text("  the particle will be on probation for "  + reqdProb + " epochs. If it hasn't improved by then,", x, y + rowHeight * i++);
      text("  the particle will be killed and respawned at the coordinates of one of the best particles.", x, y + rowHeight * i++);
    }
    i++;
    fill(255, 204, 0, 192);
    text("Normal (performing within expectations)", x, y + rowHeight * i++);
    fill(255, 102, 0, 192);
    text("On probation (liable to be terminated)", x, y + rowHeight * i++);
    fill(255, 192);
    text("Newly spawned (< 20 epochs old)", x, y + rowHeight * i++);
    fill(51, 255, 0, 192);
    text("Best particle", x, y + rowHeight * i++);
    
    //noLoop();
    drawButtons();
    return;
  }
  if (age++ < maxAge) {
    // Find the best particles overall, and add them as necessary
    bestIndices = new PriorityList(nBest);
    for (int i = 0; i < particleCount; i++) {
      particles[i].isBest = false;
      bestIndices.add(i, particles[i].fitness());
    }
    
    for (int j = 0; j < nBest; j++) {
      particles[bestIndices.items[j]].isBest = true;
    }
    
    // Slowly fade out old stuff
    noStroke();
    fill(0, 0, 0, 40);
    rect(0, 0, dimX, dimY);
    
    // Update the particles
    for (int i = 0; i < particleCount; i++) {
      // Check for improvement, and kill the particle if necessary
      if (particles[i].age - particles[i].bfage >= reqdEpochs) {
        // Check if it's improved sufficiently
        if (particles[i].bf / particles[i].fitness() >= reqdImprv) {
          if (particles[i].probation >= reqdProb) {
            // If the particle's exceeded its probation period, kill it, and respawn it from one of the best particles.
            int bIndex = int(random(nBest));
            fill(255, 255, 128, 192);
            ellipse(particles[i].x, particles[i].y, 8, 8);
            particles[i] = new Particle(i, particles[bestIndices.items[bIndex]].x, particles[bestIndices.items[bIndex]].y);
          } else {
            // Increase probation
            particles[i].probation += 2;
          }
        }
      }   
      particles[i].move();
    }
    
    // Highlight the best particles
    //stroke(255);
    //fill(255, 0, 0);
    fill(51, 255, 0, 192);
    //strokeWeight(2);
    noStroke();
    for (int j = 0; j < nBest; j++) {
      ellipse(particles[bestIndices.items[j]].x, particles[bestIndices.items[j]].y, 4, 4);
    }
    
    // Draw the target
    noStroke();
    fill(255);
    ellipse(tx, ty, 2, 2);
    
    // Draw some information
    noStroke();
    fill(0);
    rect(8, 8, 150, 60);
    fill(255);
    String s;
    s = "Target: (" + nf(tx, 2, 2) + ", " + nf(ty, 2, 2) + ")\n";
    s = s + "Epoch: " + nf(age, 3, 0);
    text(s, 12, 24);
    outputData();
  } else {
    // Highlight the best particles
    //stroke(255);
    //fill(255, 0, 0);
    fill(51, 255, 0, 192);
    //strokeWeight(2);
    noStroke();
    for (int j = 0; j < nBest; j++) {
      ellipse(particles[bestIndices.items[j]].x, particles[bestIndices.items[j]].y, 8, 8);
    }
    
    // Draw the target
    noStroke();
    fill(255);
    ellipse(tx, ty, 2, 2);
    
    // Flush the data to the file
    if (doOutputData) {
      saveStrings("output", toOutput);
    }
    noLoop();
  }
  drawButtons();
}

void outputData() {
  if (!doOutputData) return;
  float n = 0.0;
  for (int i = 0; i < particleCount; i++) {
    n += particles[i].fitness();
  }
  n = n / particleCount;
  toOutput[age] = age + "," + n;
}