#include <algorithm>
#include <cmath>
#include <cstdint>
#include <iostream>
#include <random>
#include <sophia/application.hpp>
#include <sophia/event.hpp>
#include <sophia/key_event.hpp>
#include <stdexcept>
#include <string>
#include <vector>

#define GLM_FORCE_RADIANS
#define GLM_FORCE_DEPTH_ZERO_TO_ONE
#include <glm/glm.hpp>
#include <glm/gtc/constants.hpp>
#include <glm/gtc/matrix_transform.hpp>

struct Boid {
  glm::vec2 position;
  glm::vec2 velocity;
};

constexpr uint32_t DEFAULT_NUMBER_OF_BOIDS = 25;

class Boids2D : public sophia::Application {
 public:
  static constexpr float SCROLL_SPEED = 200.0f;

  // World is a fixed box independent of the camera; boids wrap at its edges.
  static constexpr float WORLD_HALF_WIDTH = 1000.0f;
  static constexpr float WORLD_HALF_HEIGHT = 1000.0f;

  static constexpr float MAX_SPEED = 120.0f;
  static constexpr float MAX_FORCE = 200.0f;
  static constexpr float PERCEPTION_RADIUS = 150.0f;
  static constexpr float SEPARATION_RADIUS = 40.0f;
  static constexpr float ALIGNMENT_WEIGHT = 1.0f;
  static constexpr float COHESION_WEIGHT = 1.0f;
  static constexpr float SEPARATION_WEIGHT = 1.5f;

  static constexpr float BOID_LENGTH = 20.0f;
  static constexpr float BOID_WIDTH = 12.0f;

  // Neighbor lookups use a uniform grid instead of an O(n^2) scan. Cell size
  // equals PERCEPTION_RADIUS so that a boid's neighbors can only ever fall in
  // its own cell or the 8 cells surrounding it.
  static constexpr float CELL_SIZE = PERCEPTION_RADIUS;
  static constexpr int GRID_WIDTH =
      static_cast<int>(2.0f * WORLD_HALF_WIDTH / CELL_SIZE) + 1;
  static constexpr int GRID_HEIGHT =
      static_cast<int>(2.0f * WORLD_HALF_HEIGHT / CELL_SIZE) + 1;

  Boids2D() : Application{{.width = 750, .height = 1000, .name = "Boids2D"}} {
    uint32_t numberOfBoids = DEFAULT_NUMBER_OF_BOIDS;
    if (args.size() > 1) {
      numberOfBoids = static_cast<uint32_t>(std::stoul(args[1]));
    }
    boids.resize(numberOfBoids);

    sophia::EventSystem::get().addListener<sophia::KeyPressedEvent>(
        [this](sophia::KeyPressedEvent& e) {
          setScrollKeyHeld(e.getKeyCode(), true);
        });
    sophia::EventSystem::get().addListener<sophia::KeyReleasedEvent>(
        [this](sophia::KeyReleasedEvent& e) {
          setScrollKeyHeld(e.getKeyCode(), false);
        });
  }

  void onCreate() override {
    std::cout << "Boids2D created" << std::endl;

    std::mt19937 rng{std::random_device{}()};
    std::uniform_real_distribution<float> posX{-WORLD_HALF_WIDTH,
                                               WORLD_HALF_WIDTH};
    std::uniform_real_distribution<float> posY{-WORLD_HALF_HEIGHT,
                                               WORLD_HALF_HEIGHT};
    std::uniform_real_distribution<float> angle{0.0f, glm::two_pi<float>()};

    for (Boid& boid : boids) {
      boid.position = {posX(rng), posY(rng)};
      float theta = angle(rng);
      boid.velocity = glm::vec2{std::cos(theta), std::sin(theta)} * MAX_SPEED;
    }

    // Boid shape in local space, nose pointing along +x; each boid is drawn
    // as an instance of this one mesh with a per-boid transform.
    std::vector<sophia::CanvasVertex> vertices{
        {{BOID_LENGTH * 0.5f, 0.0f}, {1.0f, 1.0f, 1.0f}},
        {{-BOID_LENGTH * 0.5f, BOID_WIDTH * 0.5f}, {1.0f, 1.0f, 1.0f}},
        {{-BOID_LENGTH * 0.5f, -BOID_WIDTH * 0.5f}, {1.0f, 1.0f, 1.0f}},
    };
    boidMesh = resources().create<sophia::Mesh>(std::move(vertices));
  }

  void onUpdate(const sophia::FrameInfo& frameInfo) override {
    float scrollDistance =
        SCROLL_SPEED * static_cast<float>(frameInfo.deltaTime);
    if (scrollLeft) camera.position.x -= scrollDistance;
    if (scrollRight) camera.position.x += scrollDistance;
    if (scrollUp) camera.position.y -= scrollDistance;
    if (scrollDown) camera.position.y += scrollDistance;

    updateBoids(static_cast<float>(frameInfo.deltaTime));

    for (const Boid& boid : boids) {
      drawBoid(boid);
    }
  }

 private:
  // Maps a world position to its grid cell, clamped to the grid bounds.
  glm::ivec2 cellOf(glm::vec2 position) const {
    int cx = static_cast<int>((position.x + WORLD_HALF_WIDTH) / CELL_SIZE);
    int cy = static_cast<int>((position.y + WORLD_HALF_HEIGHT) / CELL_SIZE);
    return {std::clamp(cx, 0, GRID_WIDTH - 1),
            std::clamp(cy, 0, GRID_HEIGHT - 1)};
  }

  void rebuildGrid() {
    for (std::vector<uint32_t>& cell : grid) cell.clear();

    for (uint32_t i = 0; i < boids.size(); ++i) {
      glm::ivec2 cell = cellOf(boids[i].position);
      grid[cell.y * GRID_WIDTH + cell.x].push_back(i);
    }
  }

  void updateBoids(float deltaTime) {
    rebuildGrid();

    std::vector<glm::vec2> accelerations(boids.size());

    for (size_t i = 0; i < boids.size(); ++i) {
      glm::vec2 separation{0.0f};
      glm::vec2 averageVelocity{0.0f};
      glm::vec2 averagePosition{0.0f};
      int neighborCount = 0;

      glm::ivec2 cell = cellOf(boids[i].position);

      for (int cy = std::max(cell.y - 1, 0);
           cy <= std::min(cell.y + 1, GRID_HEIGHT - 1); ++cy) {
        for (int cx = std::max(cell.x - 1, 0);
             cx <= std::min(cell.x + 1, GRID_WIDTH - 1); ++cx) {
          for (uint32_t j : grid[cy * GRID_WIDTH + cx]) {
            if (j == i) continue;

            glm::vec2 offset = boids[i].position - boids[j].position;
            float dist = glm::length(offset);
            if (dist >= PERCEPTION_RADIUS) continue;

            averageVelocity += boids[j].velocity;
            averagePosition += boids[j].position;
            ++neighborCount;

            if (dist > 0.0f && dist < SEPARATION_RADIUS) {
              separation += offset / (dist * dist);
            }
          }
        }
      }

      glm::vec2 acceleration =
          steerTowards(boids[i], separation) * SEPARATION_WEIGHT;

      if (neighborCount > 0) {
        averageVelocity /= static_cast<float>(neighborCount);
        acceleration +=
            steerTowards(boids[i], averageVelocity) * ALIGNMENT_WEIGHT;

        averagePosition /= static_cast<float>(neighborCount);
        glm::vec2 towardsGroup = averagePosition - boids[i].position;
        acceleration += steerTowards(boids[i], towardsGroup) * COHESION_WEIGHT;
      }

      accelerations[i] = acceleration;
    }

    for (size_t i = 0; i < boids.size(); ++i) {
      Boid& boid = boids[i];
      boid.velocity += accelerations[i] * deltaTime;

      float speed = glm::length(boid.velocity);
      if (speed > MAX_SPEED) {
        boid.velocity = boid.velocity / speed * MAX_SPEED;
      }

      boid.position += boid.velocity * deltaTime;
      wrapPosition(boid.position);
    }
  }

  // Steers boid towards desiredDirection at up to MAX_SPEED, clamped to
  // MAX_FORCE.
  glm::vec2 steerTowards(const Boid& boid, glm::vec2 desiredDirection) {
    float directionLength = glm::length(desiredDirection);
    if (directionLength < 1e-5f) return glm::vec2{0.0f};

    glm::vec2 desired = desiredDirection / directionLength * MAX_SPEED;
    glm::vec2 steer = desired - boid.velocity;

    float steerLength = glm::length(steer);
    if (steerLength > MAX_FORCE) {
      steer = steer / steerLength * MAX_FORCE;
    }
    return steer;
  }

  void wrapPosition(glm::vec2& position) {
    if (position.x > WORLD_HALF_WIDTH) position.x = -WORLD_HALF_WIDTH;
    if (position.x < -WORLD_HALF_WIDTH) position.x = WORLD_HALF_WIDTH;
    if (position.y > WORLD_HALF_HEIGHT) position.y = -WORLD_HALF_HEIGHT;
    if (position.y < -WORLD_HALF_HEIGHT) position.y = WORLD_HALF_HEIGHT;
  }

  void drawBoid(const Boid& boid) {
    glm::vec2 direction = boid.velocity;
    float speed = glm::length(direction);
    direction = speed > 1e-5f ? direction / speed : glm::vec2{0.0f, -1.0f};
    float angle = std::atan2(direction.y, direction.x);

    glm::mat4 transform =
        glm::translate(glm::mat4{1.f}, glm::vec3(boid.position, 0.0f)) *
        glm::rotate(glm::mat4{1.f}, angle, glm::vec3(0.0f, 0.0f, 1.0f));

    drawMesh(boidMesh, transform);
  }

  void setScrollKeyHeld(sophia::KeyCode keycode, bool held) {
    if (keycode == sophia::Key::Left) scrollLeft = held;
    if (keycode == sophia::Key::Right) scrollRight = held;
    if (keycode == sophia::Key::Up) scrollUp = held;
    if (keycode == sophia::Key::Down) scrollDown = held;
  }

  std::vector<Boid> boids;
  std::vector<std::vector<uint32_t>> grid{
      static_cast<size_t>(GRID_WIDTH * GRID_HEIGHT)};
  sophia::MeshHandle boidMesh{};
  bool scrollLeft = false;
  bool scrollRight = false;
  bool scrollUp = false;
  bool scrollDown = false;
};

SOPHIA_MAIN(Boids2D)