#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>

#define SDA_PIN 6
#define SCL_PIN 7

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64

Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

// Number of boids
#define NUM_BOIDS 25

struct Boid {
  float x, y;
  float vx, vy;
};

Boid boids[NUM_BOIDS];

float maxSpeed = 1.5;
float neighborDist = 12;
float separationDist = 6;

void setup() {
  Wire.begin(SDA_PIN, SCL_PIN);

  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    while (1);
  }

  display.clearDisplay();

  randomSeed(analogRead(0));

  for (int i = 0; i < NUM_BOIDS; i++) {
    boids[i].x = random(SCREEN_WIDTH);
    boids[i].y = random(SCREEN_HEIGHT);
    boids[i].vx = random(-10, 10) / 10.0;
    boids[i].vy = random(-10, 10) / 10.0;
  }
}

void limitSpeed(Boid &b) {
  float speed = sqrt(b.vx * b.vx + b.vy * b.vy);
  if (speed > maxSpeed) {
    b.vx = (b.vx / speed) * maxSpeed;
    b.vy = (b.vy / speed) * maxSpeed;
  }
}

void loop() {
  for (int i = 0; i < NUM_BOIDS; i++) {
    float sepX = 0, sepY = 0;
    float alignX = 0, alignY = 0;
    float cohX = 0, cohY = 0;

    int count = 0;

    for (int j = 0; j < NUM_BOIDS; j++) {
      if (i == j) continue;

      float dx = boids[j].x - boids[i].x;
      float dy = boids[j].y - boids[i].y;
      float dist = sqrt(dx * dx + dy * dy);

      if (dist < neighborDist) {
        alignX += boids[j].vx;
        alignY += boids[j].vy;

        // Cohesion
        cohX += boids[j].x;
        cohY += boids[j].y;

        count++;

        if (dist < separationDist) {
          sepX -= dx;
          sepY -= dy;
        }
      }
    }

    if (count > 0) {
      alignX /= count;
      alignY /= count;

      cohX = (cohX / count) - boids[i].x;
      cohY = (cohY / count) - boids[i].y;
    }

    boids[i].vx += sepX * 0.05 + alignX * 0.05 + cohX * 0.01;
    boids[i].vy += sepY * 0.05 + alignY * 0.05 + cohY * 0.01;

    limitSpeed(boids[i]);
  }

  for (int i = 0; i < NUM_BOIDS; i++) {
    boids[i].x += boids[i].vx;
    boids[i].y += boids[i].vy;

    // Wrap around screen
    if (boids[i].x < 0) boids[i].x = SCREEN_WIDTH;
    if (boids[i].x >= SCREEN_WIDTH) boids[i].x = 0;
    if (boids[i].y < 0) boids[i].y = SCREEN_HEIGHT;
    if (boids[i].y >= SCREEN_HEIGHT) boids[i].y = 0;
  }

  display.clearDisplay();

  for (int i = 0; i < NUM_BOIDS; i++) {
    display.drawPixel((int)boids[i].x, (int)boids[i].y, SSD1306_WHITE);
  }

  display.display();

  delay(30);
}