/**
Click and drag to add randomized new data points to the canvas.
"SPACE" to restart cluster centers from a new location.
"UP" to increase the number of clusters.
"DOWN" to decrease the number of clusters.

*/

class Point {float x, y; Point(float x_, float y_) {x = x_; y = y_;}}
Vector points;
Point[] means;
Vector[] clusters;
Vector[] meansTrails;

int k = 5;
final int drawSize = 20;

void setup() {
  size(400,400);
  colorMode(HSB, 1.0);
  ellipseMode(CENTER_DIAMETER);
  points = new Vector();
  initMeans();
}

void initMeans() {
  if(means == null || means.length != k) {
    means = new Point[k];
    clusters = new Vector[k];
    meansTrails = new Vector[k];
    for(int i = 0; i < k; i++) {
      means[i] = randomPoint();
      clusters[i] = new Vector();
      meansTrails[i] = new Vector();
    }
  } else {
    for(int i = 0; i < k; i++) {
      means[i] = randomPoint(); 
      clusters[i].clear();
      meansTrails[i].clear();
    }
  }
  kstep();
  mstep();
}

boolean released = true;
void dokey() {
  if(keyPressed && released) {
  switch(key) {
    case UP: k++; break;
    case DOWN: k--; if(k==0) k = 1; break;
    case 8: case 127: points.clear(); break;
  }
  initMeans();
  released = false;
  }
}
void keyReleased() {
  released = true;
}

void kstep() {
  for(int i = 0; i < clusters.length; i++) clusters[i].clear();
  
  for(int i = 0; i < points.size(); i++) {
    float minDist = 10000000.0;
    int minIndex = 0;
    Point p = (Point)points.get(i);
    for(int j = 0; j < means.length; j++) {
      Point m = means[j];
      float d = dist(p.x,p.y,m.x,m.y);
      if(d < minDist) {
        minDist = d;
        minIndex = j;
      }
    }
    clusters[minIndex].add(p);
  }
}

void mstep() {
  for(int i = 0; i < clusters.length; i++) {
    float xSum = 0.0, ySum = 0.0;
    Vector cluster = clusters[i];
    for(int j = 0; j < cluster.size(); j++) {
      Point p = (Point)cluster.get(j);
      xSum += p.x;
      ySum += p.y;
    }
    if(cluster.size() > 0) {
      meansTrails[i].add(means[i]);
      means[i] = new Point(xSum/cluster.size(), ySum/cluster.size());
    } else {
      means[i] = randomPoint();
    }
  }
}

void draw() {
  background(0.0, 0.0, 0.0);
  dokey();
  if(mousePressed) {
    addPoints();
    drawPoints(); 
  } else {
    drawClusters();
    drawMeans();
  }
  kstep();
  mstep();
}

void drawClusters() {
  float colorGradation = 1/(float)clusters.length;
  for(int i = 0; i < clusters.length; i++) {
    beginShape(POINTS);
    stroke(colorGradation*i, 0.8, 0.9);
    Vector cluster = clusters[i];
    for(int j = 0; j < cluster.size(); j++) {
      Point p = (Point)cluster.get(j);
      vertex(p.x, p.y);
    }
    endShape();
  }
}

void drawMeans() {
  float colorGradation = 1/(float)clusters.length;
  for(int i = 0; i < means.length; i++) {
    Vector trail = meansTrails[i];
    stroke(0,0,0.5);
    noFill();
    beginShape();
    for(int j = 0; j < trail.size(); j++) {
      Point p = (Point)trail.get(j);
      vertex(p.x, p.y);
    }
    endShape();
    fill(colorGradation*i, 0.8, 0.9);
    stroke(0,0,1.0);
    ellipse(means[i].x, means[i].y, 4, 4);
  }
}

void drawPoints() {
  beginShape(POINTS);
  stroke(0,0,1.0);
  for(int i = 0; i < points.size(); i++) {
    Point p = (Point)points.get(i);
    vertex(p.x, p.y);
  }
  endShape();
}

void addPoints() {
  points.add(new Point(mouseX+random(-drawSize,  drawSize),   mouseY+random(-drawSize,  drawSize)));
  points.add(new Point(mouseX+random(-drawSize/2,drawSize/2), mouseY+random(-drawSize/2,drawSize/2)));
  points.add(new Point(mouseX+random(-drawSize/4,drawSize/4), mouseY+random(-drawSize/4,drawSize/4)));
}

Point randomPoint() { return new Point(random(width), random(height)); }