#include <SDL2/SDL.h>
#include <math.h>
#include <stdlib.h>
#include <time.h>

#define SCREEN_WIDTH 800
#define SCREEN_HEIGHT 600
#define NUM_PARTICLES 200
#define NUM_TYPES 2
#define FORCE_SCALE 100.0
#define PARTICLE_RADIUS 5
#define GRID_SIZE 50 // Spatial grid size

typedef struct Particle {
  float x, y;            // Position
  float vx, vy;          // Velocity
  int type;              // Particle type (0, 1, 2... NUM_TYPES-1)
  struct Particle *next; // Linked list pointer for spatial grid
} Particle;

typedef struct {
  int x, y; // Top-left corner of slider
  int width, height;
  float *value; // Pointer to the value it modifies
} Slider;

Particle particles[NUM_PARTICLES];
float force_matrix[NUM_TYPES][NUM_TYPES];
Slider sliders[NUM_TYPES * NUM_TYPES];
int grid_cols, grid_rows;
Particle **grid = NULL; // Dynamic grid allocation

void init_particles() {
  srand((unsigned int)time(NULL));
  for (int i = 0; i < NUM_PARTICLES; i++) {
    particles[i].x = rand() % SCREEN_WIDTH;
    particles[i].y = rand() % SCREEN_HEIGHT;
    particles[i].vx = (rand() % 200 - 100) / 100.0f;
    particles[i].vy = (rand() % 200 - 100) / 100.0f;
    particles[i].type = rand() % NUM_TYPES;
    particles[i].next = NULL;
  }
}

void init_force_matrix() {
  for (int i = 0; i < NUM_TYPES; i++) {
    for (int j = 0; j < NUM_TYPES; j++) {
      force_matrix[i][j] = (rand() % 200 - 100) * FORCE_SCALE / 100.0f;
    }
  }
}

void init_sliders() {
  int slider_width = 100;
  int slider_height = 10;
  int spacing = 20;
  int start_x = 10, start_y = SCREEN_HEIGHT - NUM_TYPES * spacing - 30;

  for (int i = 0; i < NUM_TYPES; i++) {
    for (int j = 0; j < NUM_TYPES; j++) {
      Slider *slider = &sliders[i * NUM_TYPES + j];
      slider->x = start_x + j * (slider_width + 5);
      slider->y = start_y + i * spacing;
      slider->width = slider_width;
      slider->height = slider_height;
      slider->value = &force_matrix[i][j];
    }
  }
}

void handle_sliders(SDL_Event *event) {
  if (event->type == SDL_MOUSEBUTTONDOWN || event->type == SDL_MOUSEMOTION) {
    int mx, my;
    Uint32 buttons = SDL_GetMouseState(&mx, &my);
    if (buttons & SDL_BUTTON(SDL_BUTTON_LEFT)) {
      for (int i = 0; i < NUM_TYPES * NUM_TYPES; i++) {
        Slider *slider = &sliders[i];
        if (mx >= slider->x && mx <= slider->x + slider->width &&
            my >= slider->y && my <= slider->y + slider->height) {
          float new_value = (float)(mx - slider->x) / slider->width * 2 - 1;
          *slider->value = new_value * FORCE_SCALE;
        }
      }
    }
  }
}

void update_particles(float dt) {
  // Clear the grid
  for (int i = 0; i < grid_cols * grid_rows; i++) {
    grid[i] = NULL;
  }

  // Assign particles to the grid
  for (int i = 0; i < NUM_PARTICLES; i++) {
    int gx = particles[i].x / GRID_SIZE;
    int gy = particles[i].y / GRID_SIZE;
    if (gx >= 0 && gx < grid_cols && gy >= 0 && gy < grid_rows) {
      int cell_index = gy * grid_cols + gx;
      particles[i].next = grid[cell_index];
      grid[cell_index] = &particles[i];
    }
  }

  // Update particles
  for (int i = 0; i < NUM_PARTICLES; i++) {
    Particle *p = &particles[i];
    int gx = p->x / GRID_SIZE;
    int gy = p->y / GRID_SIZE;

    float fx = 0, fy = 0;

    // Check surrounding grid cells
    for (int dx = -1; dx <= 1; dx++) {
      for (int dy = -1; dy <= 1; dy++) {
        int ngx = gx + dx;
        int ngy = gy + dy;
        if (ngx >= 0 && ngx < grid_cols && ngy >= 0 && ngy < grid_rows) {
          int cell_index = ngy * grid_cols + ngx;
          Particle *neighbor = grid[cell_index];
          while (neighbor) {
            if (neighbor != p) {
              float dx = neighbor->x - p->x;
              float dy = neighbor->y - p->y;
              float dist2 = dx * dx + dy * dy;
              if (dist2 > 0 && dist2 < 2500) { // Squared distance
                float force =
                    force_matrix[p->type][neighbor->type] / sqrtf(dist2);
                fx += force * dx;
                fy += force * dy;
              }
            }
            neighbor = neighbor->next;
          }
        }
      }
    }

    // Apply forces and update positions
    p->vx += fx * dt;
    p->vy += fy * dt;
    p->x += p->vx * dt;
    p->y += p->vy * dt;

    // Handle boundaries
    if (p->x < 0 || p->x > SCREEN_WIDTH)
      p->vx *= -1;
    if (p->y < 0 || p->y > SCREEN_HEIGHT)
      p->vy *= -1;
    p->x = fmaxf(0, fminf(SCREEN_WIDTH, p->x));
    p->y = fmaxf(0, fminf(SCREEN_HEIGHT, p->y));
  }
}

void render_particles(SDL_Renderer *renderer) {
  for (int i = 0; i < NUM_PARTICLES; i++) {
    Particle *p = &particles[i];
    switch (p->type) {
    case 0:
      SDL_SetRenderDrawColor(renderer, 255, 0, 0, 255);
      break;
    case 1:
      SDL_SetRenderDrawColor(renderer, 0, 255, 0, 255);
      break;
    case 2:
      SDL_SetRenderDrawColor(renderer, 0, 0, 255, 255);
      break;
    default:
      SDL_SetRenderDrawColor(renderer, 255, 255, 255, 255);
      break;
    }
    SDL_Rect rect = {(int)p->x - PARTICLE_RADIUS, (int)p->y - PARTICLE_RADIUS,
                     PARTICLE_RADIUS * 2, PARTICLE_RADIUS * 2};
    SDL_RenderFillRect(renderer, &rect);
  }
}

void render_sliders(SDL_Renderer *renderer) {
  for (int i = 0; i < NUM_TYPES * NUM_TYPES; i++) {
    Slider *slider = &sliders[i];
    SDL_SetRenderDrawColor(renderer, 200, 200, 200, 255);
    SDL_Rect rect = {slider->x, slider->y, slider->width, slider->height};
    SDL_RenderFillRect(renderer, &rect);

    SDL_SetRenderDrawColor(renderer, 0, 255, 0, 255);
    int handle_x =
        slider->x + ((*slider->value / FORCE_SCALE + 1) / 2) * slider->width;
    SDL_Rect handle = {handle_x - 5, slider->y, 10, slider->height};
    SDL_RenderFillRect(renderer, &handle);
  }
}

int main(int argc, char *argv[]) {
  if (SDL_Init(SDL_INIT_VIDEO) < 0) {
    fprintf(stderr, "Could not initialize SDL: %s\n", SDL_GetError());
    return 1;
  }

  SDL_Window *window = SDL_CreateWindow(
      "Particle Simulator", SDL_WINDOWPOS_CENTERED, SDL_WINDOWPOS_CENTERED,
      SCREEN_WIDTH, SCREEN_HEIGHT, SDL_WINDOW_SHOWN);
  if (!window) {
    fprintf(stderr, "Could not create window: %s\n", SDL_GetError());
    SDL_Quit();
    return 1;
  }

  SDL_Renderer *renderer =
      SDL_CreateRenderer(window, -1, SDL_RENDERER_ACCELERATED);
  if (!renderer) {
    fprintf(stderr, "Could not create renderer: %s\n", SDL_GetError());
    SDL_DestroyWindow(window);
    SDL_Quit();
    return 1;
  }

  // Initialize components
  init_particles();
  init_force_matrix();
  init_sliders();

  grid_cols = SCREEN_WIDTH / GRID_SIZE;
  grid_rows = SCREEN_HEIGHT / GRID_SIZE;
  grid = (Particle **)calloc(grid_cols * grid_rows, sizeof(Particle *));

  int running = 1;
  SDL_Event event;
  Uint32 last_time = SDL_GetTicks();

  while (running) {
    while (SDL_PollEvent(&event)) {
      if (event.type == SDL_QUIT) {
        running = 0;
      } else {
        handle_sliders(&event);
      }
    }

    Uint32 current_time = SDL_GetTicks();
    float dt = (current_time - last_time) / 1000.0f;
    last_time = current_time;

    update_particles(dt);

    SDL_SetRenderDrawColor(renderer, 0, 0, 0, 255);
    SDL_RenderClear(renderer);
    render_particles(renderer);
    render_sliders(renderer);
    SDL_RenderPresent(renderer);
  }

  // Clean up
  free(grid);
  SDL_DestroyRenderer(renderer);
  SDL_DestroyWindow(window);
  SDL_Quit();
  return 0;
}