#include <cmath>
#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>

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

constexpr uint32_t DEFAULT_NUMBER_OF_BOIDS = 25;

class Testbed : 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 = 375.0f;
  static constexpr float WORLD_HALF_HEIGHT = 500.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;

  Testbed() : Application{{.width = 750, .height = 1000, .name = "Testbed"}} {
    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 << "Testbed 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;
    }
  }

  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));

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

 private:
  void updateBoids(float deltaTime) {
    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;

      for (size_t j = 0; j < boids.size(); ++j) {
        if (i == j) 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) {
    constexpr float LENGTH = 20.0f;
    constexpr float WIDTH = 12.0f;

    glm::vec2 direction = boid.velocity;
    float speed = glm::length(direction);
    direction = speed > 1e-5f ? direction / speed : glm::vec2{0.0f, -1.0f};
    glm::vec2 perpendicular{-direction.y, direction.x};

    glm::vec2 nose = boid.position + direction * (LENGTH * 0.5f);
    glm::vec2 tailLeft = boid.position - direction * (LENGTH * 0.5f) +
                         perpendicular * (WIDTH * 0.5f);
    glm::vec2 tailRight = boid.position - direction * (LENGTH * 0.5f) -
                          perpendicular * (WIDTH * 0.5f);

    canvas.triangle(nose, tailLeft, tailRight, {1.0f, 1.0f, 1.0f});
  }

  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;
  bool scrollLeft = false;
  bool scrollRight = false;
  bool scrollUp = false;
  bool scrollDown = false;
};

SOPHIA_MAIN(Testbed)