antcolony.pde

int WIDTH = 500;
int HEIGHT = 500;
int ANTS_AT_A_TIME = 30;
boolean DRAW_PATH = false;

float sum_array( float[] floats ) {
  float ret = 0;
  
  for(int i=0; i<floats.length; i++){
    ret += floats[i];
  }
  
  return ret;
}

int choose( float[] classes ) {
  float classes_sum = sum_array( classes );
  
  float spinner_value = random( classes_sum );
  
  // find which class the spinner landed on
  float curr_spinner_upper_bound = 0;
  for(int i=0; i<classes.length; i++) {
    curr_spinner_upper_bound += classes[i];
    
    if(curr_spinner_upper_bound > spinner_value) {
      return i;
    }
  }
  
  //this really should never happen
  return classes.length-1;
}

class Ant {
  int SEARCH_MODE=0;
  int HOME_MODE=1;
  
  int home_x;
  int home_y;
  int x;
  int y;
  int mode;
  
  int[] path;
  int path_head;
  
  Ant( int x, int y ) {
    this.home_x = x;
    this.home_y = y;
    this.x = x;
    this.y = y;
    
    this.path = new int[5000];
    this.path[0] = y*WIDTH+x;
    this.path_head=1;
    
    this.mode = SEARCH_MODE;
  }
  
  void draw() {
    point( this.x, this.y ); 
  }
  
  void draw_path() {
    for(int i=0; i<this.path_head; i++){
      int y = this.path[i]/WIDTH;
      int x = this.path[i]%WIDTH;
      point(x,y);
    }
  }
  
  void update_path(int x, int y) {
    int path_head_contents = y*WIDTH+x;
    
    for(int i=0; i<path.length; i++) {
      if( path[i] == path_head_contents ) {
        this.path_head = i;
        return;
      }
    }
    
    this.path[path_head] = path_head_contents;
    this.path_head += 1;
  }
  
  void move() {
    if( this.mode == SEARCH_MODE ) {
      wander();
      
      if( food[y*WIDTH+x] ) {
        this.mode = HOME_MODE;
      }
    } else if( this.mode == HOME_MODE ) {
      backtrack();
    }
  }
  
  void wander() {
    
    float[] chances = { 1+pherimone[(y-1)*WIDTH+x], 1+pherimone[y*WIDTH+x+1], 1+pherimone[(y+1)*WIDTH+x], 1+pherimone[y*WIDTH+x-1] };
    int direction = choose( chances );
    
    if(direction==0){
      this.y -= 1;
    }else if(direction==1){
      this.x += 1;
    }else if(direction==2){
      this.y += 1;
    }else if(direction==3){
      this.x -= 1;
    }
    
    this.update_path(this.x, this.y);
  }
  
  void deposit_pherimone() {
    pherimone[y*WIDTH+x] = constrain(pherimone[y*WIDTH+x]+2,0,255);
    pherimone_image.set( x, y, color(255-pherimone[y*WIDTH+x]) );
  }
  
  void backtrack() {
    if(this.path_head!=0){
      deposit_pherimone();
      
      this.path_head -= 1;
    
      this.y = this.path[this.path_head]/WIDTH;
      this.x = this.path[this.path_head]%WIDTH;
    } else {
      this.mode = SEARCH_MODE;
    }
  }
}

int ANT_MODE = 0;
int FOOD_MODE = 1;

Vector ants;
boolean[] food;
int[] pherimone;
PImage pherimone_image;
int mode;

void setup() {
  size(WIDTH, HEIGHT);
  smooth();
  strokeWeight(2);
  ants = new Vector();
  food = new boolean[WIDTH*HEIGHT];
  
  pherimone = new int[WIDTH*HEIGHT];
  for(int i=0; i<WIDTH*HEIGHT; i++) {
    pherimone[i] = 0; 
  }
  
  pherimone_image = createImage( WIDTH, HEIGHT, RGB );
  pherimone_image.loadPixels();
  for (int i = 0; i < pherimone_image.pixels.length; i++) {
   pherimone_image.pixels[i] = color(255); 
  }
  pherimone_image.updatePixels();
}

void draw() {
  image( pherimone_image, 0, 0 );
  
  strokeWeight( 2 );
  stroke(0);
  for( int i=0; i<ants.size(); i++) {
    Ant ant = (Ant)ants.get( i );
    ant.draw();
    ant.move();
  }
  
  if(DRAW_PATH){
    strokeWeight(1);
    stroke(0,255,0);
    for( int i=0; i<ants.size(); i++) {
      Ant ant = (Ant)ants.get( i );
      ant.draw_path();
    }
  }
  
  strokeWeight(1);
  stroke(255,0,0);
  for(int y=0; y<HEIGHT; y++) {
    for(int x=0; x<WIDTH; x++) {
      if( food[WIDTH*y+x] ) {
        point( x, y );
      }
    }
  }
  
  
  
  /*for(int y=0; y<HEIGHT; y++) {
    for(int x=0; x<WIDTH; x++) {
      stroke( 255-pherimone[WIDTH*y+x] );
      point( x, y );
    }
  }*/
}

void keyPressed() {
  if( key == 'a' ) {
    mode = ANT_MODE;
    println( "ant mode" );
  } else if( key == 'f' ) {
    mode = FOOD_MODE;
    println( "food mode" );
  }
}

void mousePressed() {
  if( mode == ANT_MODE ) {
    for(int i=0; i<ANTS_AT_A_TIME; i++) {
      Ant ant = new Ant( mouseX, mouseY );
      ants.add( ant );
    }
  } else if( mode == FOOD_MODE ) {
    food[WIDTH*mouseY+mouseX] = true;
  }
}

void mouseDragged() {
  if( mode == FOOD_MODE ) {
    food[WIDTH*mouseY+mouseX] = true;
  }
}