#include "board.hpp"
#include "constants.hpp"
#include "forces.hpp"
#include "pedestrian.hpp"
#include "social_force.hpp"
#include "vector2d.hpp"

#include <algorithm>
#include <cmath>
#include <fstream>
#include <iostream>
#include <random>
#include <tuple>
#include <vector>

#define SEED 42

#define BOARD_WIDTH 10.0
#define BOARD_LENGTH 50.0
#define WALL_EPS 1e-3
#define BOARD_WIDTH_MINUS_WALL_EPS 9.999

#ifndef RANDOM_SEED
// Fixed seeding is used by default to make simulations reproducible.
#define RANDOM_SEED 0
#endif

#ifndef BENCHMARK
// Benchmarking is disabled by default.
#define BENCHMARK 0
#endif

using namespace std;

Pedestrian *init_pedestrians(int n_p) {
#if RANDOM_SEED
  random_device rd;  // Platform-provided random seed source.
  mt19937 gen(rd()); // Non-reproducible initialization.
#else
  mt19937 gen(SEED); // Reproducible initialization.
#endif

  normal_distribution<double> velocity_gen(v0_mean, v0_std_dev);
  uniform_real_distribution<double> y_gen(WALL_EPS, BOARD_WIDTH_MINUS_WALL_EPS);

  Pedestrian *P = new Pedestrian[n_p];

  int n_p_half = n_p / 2;

  // Group of pedestrians walking from left to right.
  for (int i = 0; i < n_p_half; i++) {
    double x = 0.0;
    double y = y_gen(gen);

    Vector2d start_pos = Vector2d(x, y);
    Vector2d destination = Vector2d(BOARD_LENGTH, y);

    Vector2d init_e = destination - start_pos;
    double init_e_norm = init_e.norm();
    Vector2d init_e_unit = init_e / init_e_norm;

    // Reject near-zero desired speeds to avoid pedestrians that barely move,
    // and reject negative speeds because they would reverse the intended
    // walking direction.
    double desired_v0;
    do {
      desired_v0 = velocity_gen(gen);
    } while (desired_v0 <= 0.1);

    double v_max = 1.3 * desired_v0;

    P[i] = Pedestrian(WHITE, Vector2d(), start_pos, init_e_unit, v_max, desired_v0, destination);
  }

  // Group of pedestrians walking from right to left.
  for (int i = n_p_half; i < n_p; i++) {
    double x = BOARD_LENGTH;
    double y = y_gen(gen);

    Vector2d start_pos = Vector2d(x, y);
    Vector2d destination = Vector2d(0.0, y);

    Vector2d init_e = destination - start_pos;
    double init_e_norm = init_e.norm();
    Vector2d init_e_unit = init_e / init_e_norm;

    // Reject near-zero desired speeds to avoid pedestrians that barely move,
    // and reject negative speeds because they would reverse the intended
    // walking direction.
    double desired_v0;
    do {
      desired_v0 = velocity_gen(gen);
    } while (desired_v0 <= 0.1);

    double v_max = 1.3 * desired_v0;

    P[i] = Pedestrian(BLACK, Vector2d(), start_pos, init_e_unit, v_max, desired_v0, destination);
  }

  return P;
}

Board init_board_without_obstacles() { return Board(BOARD_LENGTH, BOARD_WIDTH, 0, nullptr); }

int run_simulation(Pedestrian *pedestrians, Board &board, int n_p, int n_iter, const char *dump_path) {

#if BENCHMARK
  // Avoid compiler unused parameter warning
  (void)dump_path;
#endif

#if !BENCHMARK
  const int dump_interval = 1;

  // Each sample stores: time step, pedestrian id, x, y, e.x, e.y, color.
  vector<tuple<int, int, double, double, double, double, int>> samples;

  if (dump_path) {
    samples.reserve(static_cast<size_t>(n_iter / dump_interval + 1) * n_p);

    // Store the initial state before the first simulation step.
    for (int i = 0; i < n_p; i++) {
      Pedestrian &p = pedestrians[i];
      samples.emplace_back(0, i, p.curr_pos.x, p.curr_pos.y, p.e.x, p.e.y, p.color);
    }
  }
#endif

  for (int t = 0; t < n_iter; t++) {
    // Pass 1: Compute all forces from the current state.
    for (int i = 0; i < n_p; i++) {
      Pedestrian &p = pedestrians[i];

      // # 2 ADDs, 4 MULTs
      //
      // $\vec{F}_\alpha(\vec{v}_\alpha, v^0_\alpha \vec{e}_\alpha) = \frac{1}{\tau_\alpha}(v^0_\alpha \vec{e}_\alpha - \vec{v}_\alpha)$
      //
      Vector2d f_alpha_dt = acceleration_term_dt(p);

      // # (N-1) * (12 ADDs, 23 MULTs, 1 DIV, 4 SQRTs, 1 EXP)
      //   + (N-1) * 2 ADDs
      //
      // $\sum_{\beta} \vec{F}_{\alpha\beta}(\vec{e}_{\alpha},\vec{r}_{\alpha} - \vec{r}_{\beta})$
      //
      for (int j = 0; j < n_p; j++) {
        if (i == j)
          continue;

        f_alpha_dt += pedestrian_repulsion_dt(p, pedestrians[j]);
      }

      // # board.n_obstacles * (10 ADDs, 18 MULTs, 1 DIV, 2 SQRTs, 1 EXP)
      //   + board.n_obstacles * 2 ADDs
      //
      // $\sum_{B} \vec{F}_{\alpha B}(\vec{e}_{\alpha}, \vec{r}_{\alpha} - \vec{r}_{B})$
      //
      for (int k = 0; k < board.n_obstacles; k++) {
        f_alpha_dt += obstacle_repulsion_dt(p, board.obstacles[k]);
      }

      // # 2 ADDs
      //
      //
      // $\frac{d\vec{w}_\alpha}{dt} = \vec{F}_{\alpha}(t)$
      //
      p.w += f_alpha_dt;
    }

    // Pass 2: Update pedestrian states after all forces have been computed.
    // # (7 ADDs, 16 MULTs, 2 DIVs, 3 SQRTs) * n_p
    for (int i = 0; i < n_p; i++) {
      Pedestrian &p = pedestrians[i];

      // # 1 ADD, 2 MULTs, 1 SQRT
      double w_norm = p.w.norm();
      double w_norm2 = w_norm > EPS ? w_norm : EPS; // Avoid division by zero.

      // # 1 DIV
      double correction_factor = std::min(1.0, p.v_max / w_norm2);

      // # 2 MULTs
      p.curr_vel = p.w * correction_factor;

      // # 2 ADDs, 2 MULTs
      //
      // $\vec{r}_i^{t+1} = \vec{r}_i^t + \vec{v}_i^t dt$
      //
      p.curr_pos = p.curr_pos + p.curr_vel * dt;

      // When a pedestrian reaches the end of the corridor, teleport them back
      // to the opposite entrance while keeping the same y-coordinate. This
      // maintains a constant pedestrian flow and makes lane formation easier to
      // observe.
      if (p.color == WHITE && p.curr_pos.x >= BOARD_LENGTH) {
        p.curr_pos.x = 0.0;
        p.w = Vector2d();
        p.curr_vel = Vector2d();
      }
      if (p.color == BLACK && p.curr_pos.x <= 0.0) {
        p.curr_pos.x = BOARD_LENGTH;
        p.w = Vector2d();
        p.curr_vel = Vector2d();
      }

      // Keep pedestrians inside the board boundaries.
      p.curr_pos.y = clamp(p.curr_pos.y, WALL_EPS, BOARD_WIDTH_MINUS_WALL_EPS);

      // # 2 ADDs
      Vector2d e = p.destination - p.curr_pos;

      // # 1 ADD, 2 MULTs, 1 DIV, 1 SQRT
      double e_norm = e.norm();
      double e_norm2 = e_norm > EPS ? e_norm : EPS; // Avoid division by zero.
      double e_norm_inv = 1.0 / e_norm2;

      // # 2 MULTs
      p.e = e * e_norm_inv;

      // # 1 ADD, 2 MULTs
      double velocity_norm_sq = p.curr_vel.dot(p.curr_vel);

      // # 1 SQRT
      double velocity_norm = std::sqrt(velocity_norm_sq);

      // # 1 MULT
      p.s_b_sq = velocity_norm_sq * delta_t_sq;

      // # 3 MULTs
      p.s_b_e = p.e * (velocity_norm * delta_t);

#if !BENCHMARK
      // Store the current pedestrian state for visualization.
      if (dump_path && (t % dump_interval == 0)) {
        samples.emplace_back(t + 1, i, p.curr_pos.x, p.curr_pos.y, p.e.x, p.e.y, p.color);
      }
#endif
    }
  }

#if !BENCHMARK
  // Export recorded simulation states to CSV.
  if (dump_path) {
    ofstream out(dump_path);
    out << "time,id,x,y,p.e.x,p.e.y,color\n";

    for (auto &[t, id, x, y, e_x, e_y, c] : samples) {
      out << t << ',' << id << ',' << x << ',' << y << ',' << e_x << ',' << e_y << ',' << c << '\n';
    }
  }
#endif

  return 0;
}