#include "./headers/board.hpp"
#include "./headers/constants.hpp"
#include "./headers/pedestrian.hpp"
#include "./headers/social_force.hpp"

#include <algorithm>
#include <cstdio>
#include <cstdlib>
#include <vector>

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

#ifndef ROOFLINE
#define ROOFLINE 0
#endif

#ifdef __x86_64__
#include "./headers/tsc_x86.h"
#include <immintrin.h>
#define TRASH_BUFFER_SIZE (16ULL * 1024 * 1024 * 2) // Size of L3 cache : 16MB
#define PAGE_SIZE 4096
#endif

#ifdef __aarch64__
#include "./headers/kperf.h"
#include <arm_acle.h>
#define TRASH_BUFFER_SIZE (1502838784ULL * 2) // Size of SLC on M3 pro chips : 1.5GB
#define PAGE_SIZE 16384
#endif

#define N_OBSTACLES 2
#define WARMUP_ITERS 5
#define NUM_REPS 11

// Flushes caches by streaming reads through a large persistent buffer.
long long flush_cache() {
  static volatile char *trash_buffer = nullptr;
  static bool initialized = false;

  if (!initialized) {
    trash_buffer = new volatile char[TRASH_BUFFER_SIZE];

    // Touch one byte per page once so pages are physically mapped. And we avoid paying the cost of page faults during cache trashing below.
    for (size_t i = 0; i < TRASH_BUFFER_SIZE; i += PAGE_SIZE) {
      trash_buffer[i] = 1;
    }

    initialized = true;
  }

  long long acc = 0;

  for (size_t i = 0; i < TRASH_BUFFER_SIZE; i += CACHE_BLOCK_SIZE) {
    // Access one byte per cache block is enough
    acc += trash_buffer[i];
  }

#ifdef __x86_64__
  _mm_lfence();
#endif

#ifdef __aarch64__
  __isb(0xF);
#endif

  asm volatile("" ::: "memory");

  return acc;
}

int main(int argc, const char *argv[]) {
  // Check that the correct number of arguments has been passed.
  //
  // Expected usage:
  //
  //   ./main <number_of_pedestrians> <number_of_iterations> [csv_output_path]
  //
  // The optional third argument enables CSV dumping for visualization.
  if (argc < 3) {
    printf("Usage: ./main <number_of_pedestrians> <number_of_iterations> [csv_output_path]\n");
    printf("Each iteration takes %.2f seconds in the simulation.\n", dt);
    return 1;
  }

  // Parse command-line arguments.
  int n_p = atoi(argv[1]);
  int n_iter = atoi(argv[2]);

  if (n_p <= 0) {
    printf("Error: <number_of_pedestrians> must be a positive integer, got '%s' instead\n", argv[1]);
    return 1;
  }

  if (n_iter <= 0) {
    printf("Error: <number_of_iterations> must be a positive integer, got '%s' instead\n", argv[2]);
    return 1;
  }

  // If provided, dump_path is used to write simulation states to CSV.
  // In benchmark mode, this path is intentionally ignored to avoid measuring
  // file I/O overhead.
  const char *dump_path = (argc >= 4) ? argv[3] : nullptr;

#if BENCHMARK

#if ROOFLINE
  volatile long long trash_value;
#endif

#ifdef __aarch64__
  // Initialize ARM performance counters.
  kperf_init();
  struct performance_counters startperf, endperf;
#endif

  // Warmup runs are not measured.
  // Each run gets a fresh initialized state because run_simulation mutates the
  // pedestrians and board state.
  for (int i = 0; i < WARMUP_ITERS; i++) {
    Pedestrian *P = init_pedestrians(n_p);
    Board board = init_board_without_obstacles();

    run_simulation(P, board, n_p, n_iter, nullptr);

    delete[] P;
  }

  // Store the cycle count of each measured repetition.
  // The median is then used as the final timing result.
  std::vector<long double> cycles_v;
  cycles_v.reserve(NUM_REPS);

  for (int i = 0; i < NUM_REPS; i++) {
    Pedestrian *P = init_pedestrians(n_p);
    Board board = init_board_without_obstacles();
#if ROOFLINE
    trash_value = flush_cache();
#endif

#ifdef __x86_64__
    myInt64 start = start_tsc();
#endif

#ifdef __aarch64__
    startperf = kperf_get_counters();
#endif

    // Disable CSV dumping during benchmarking to avoid measuring file I/O
    // overhead. Initialization is also excluded from the measured region.
    run_simulation(P, board, n_p, n_iter, nullptr);

#ifdef __x86_64__
    long double cycles = stop_tsc(start);
#endif

#ifdef __aarch64__
    endperf = kperf_get_counters();
    long double cycles = endperf.cycles - startperf.cycles;
#endif

    delete[] P;

    cycles_v.push_back(cycles);
  }

  // Use the median cycle count to reduce the effect of noisy measurements.
  std::sort(cycles_v.begin(), cycles_v.end());
  long double T = cycles_v[NUM_REPS / 2];

#else

  // Normal execution mode. If dump_path is provided, the simulation writes a CSV
  // file that can later be used to generate a video of the simulation.
  Pedestrian *P = init_pedestrians(n_p);
  Board board = init_board_without_obstacles();
  run_simulation(P, board, n_p, n_iter, dump_path);

  delete[] P;

#endif

#if BENCHMARK

  // Flop-count per simulation iteration.
  //
  // The total work W(n) is estimated from the number of additions,
  // multiplications, divisions, square roots, and exponentials performed by the
  // simulation kernels.
  long long add_count_per_iter = n_p * (4 + (n_p - 1) * 14 + N_OBSTACLES * 12) + n_p * 7;

  long long mult_count_per_iter = n_p * (4 + (n_p - 1) * 23 + N_OBSTACLES * 18) + n_p * 16;

  long long div_count_per_iter = n_p * ((n_p - 1) * 1 + N_OBSTACLES * 1) + n_p * 2;

  long long sqrt_count_per_iter = n_p * ((n_p - 1) * 4 + N_OBSTACLES * 2) + n_p * 3;

  long long exp_count_per_iter = n_p * ((n_p - 1) * 1 + N_OBSTACLES * 1);

  long double W_per_iter = add_count_per_iter * cost_add + mult_count_per_iter * cost_mult + div_count_per_iter * cost_div + sqrt_count_per_iter * cost_sqrt + exp_count_per_iter * cost_exp;

  // Estimated memory traffic per iteration.
  // TODO: refine this model to account for actual memory accesses, cache reuse,
  // and the full interaction loop.
  long long Q_per_iter = sizeof(Pedestrian) * n_p + sizeof(Obstacle) * N_OBSTACLES;

  // Operational intensity:
  //
  //   I(n) = W(n) / Q(n)
  //
  // where W is estimated work in flops and Q is estimated memory traffic in
  // bytes.
  long double I = W_per_iter / Q_per_iter;

  // Total estimated work over all iterations.
  long double W = W_per_iter * n_iter;

  // Performance in estimated flops per cycle.
  long double P = W / T;

  printf("W(n) per iteration: %.2Lf flops\n", W_per_iter);
  printf("Q(n) per iteration: %lld bytes\n", Q_per_iter);
  printf("I(n): %.2Lf flops/byte\n", I);
  printf("W(n): %.2Lf flops\n", W);
  printf("T(n): %.2Lf cycles\n", T);
  printf("P(n): %.2Lf flops/cycle\n", P);

#endif

  return 0;
}