#include "sims.h"
#include "utils/functions.h"
#include "utils/io.h"

#include <chrono>
#include <iostream>
#include <omp.h>

void run_simulation(Experiment &experiment,
                    const std::unordered_map<std::string, double> &parameter,
                    const std::unordered_map<std::string, double> &optional,
                    Dynamics &dynamics, std::mt19937_64 &rng) {
  const int num_walker = static_cast<int>(parameter.at("num_walker"));

  dynamics.setParameters(parameter);
  experiment.setup(parameter, optional);

  const auto start = printStart(optional);

  const int num_threads = omp_get_max_threads();

  // Create per-thread RNGs seeded deterministically from the main RNG
  std::vector<std::mt19937_64> thread_rngs;
  thread_rngs.reserve(num_threads);
  for (int i = 0; i < num_threads; i++) {
    thread_rngs.emplace_back(rng());
  }

  // Create per-thread clones of dynamics and experiment
  std::vector<std::unique_ptr<Dynamics>> thread_dynamics;
  std::vector<std::unique_ptr<Experiment>> thread_experiments;
  for (int i = 0; i < num_threads; i++) {
    thread_dynamics.push_back(dynamics.clone());
    thread_experiments.push_back(experiment.clone());
  }

  int steps_done = 0;
  auto last_print_out = std::chrono::system_clock::now();

#pragma omp parallel
  {
    const int tid = omp_get_thread_num();
    auto &local_rng = thread_rngs[tid];
    auto &local_dynamics = *thread_dynamics[tid];
    auto &local_experiment = *thread_experiments[tid];

#pragma omp for schedule(static)
    for (int mol_i = 0; mol_i < num_walker; mol_i++) {
      auto traj = make_trajectory(local_dynamics, experiment.tmax(), local_rng);
      local_experiment.accumulate(traj, local_dynamics.getInitOmega(),
                                  num_walker);

      if (tid == 0) {
#pragma omp atomic
        steps_done++;
        last_print_out =
            printSteps(last_print_out, start, num_walker, steps_done);
      } else {
#pragma omp atomic
        steps_done++;
      }
    }
  }

  // Merge per-thread results
  for (int i = 0; i < num_threads; i++) {
    experiment.merge(*thread_experiments[i]);
  }

  const auto directory = make_directory(dynamics.toString());
  experiment.save(directory);
  printEnd(start);
}

Trajectory make_trajectory(Dynamics &dynamics, const double t_max,
                           std::mt19937_64 &rng) {
  double t_passed = 0;
  double phase = 0;

  dynamics.initialize(rng);

  double omega = dynamics.getInitOmega();

  Trajectory traj;
  traj.time.emplace_back(t_passed);
  traj.phase.emplace_back(phase);
  traj.omega.emplace_back(omega);

  while (t_passed < t_max) {
    auto [dt, new_omega] = dynamics.next(rng);
    phase += omega * dt;
    t_passed += dt;
    omega = new_omega;

    traj.time.emplace_back(t_passed);
    traj.phase.emplace_back(phase);
    traj.omega.emplace_back(omega);
  }

  return traj;
}

std::chrono::system_clock::time_point
printStart(const std::unordered_map<std::string, double> &optional) {
  const auto start = std::chrono::system_clock::now();
  const time_t start_time = std::chrono::system_clock::to_time_t(start);

  std::cout << "Random walk for ";
  for (const auto &[key, value] : optional) {
    std::cout << key << " = " << value << "; ";
  }
  std::cout << std::endl;
  std::cout << "Start: " << ctime(&start_time);

  return start;
}

void printEnd(const std::chrono::system_clock::time_point start) {
  const auto end = std::chrono::system_clock::now();

  const std::chrono::duration<float> duration = end - start;
  const time_t end_time = std::chrono::system_clock::to_time_t(end);
  std::cout << "End:   " << ctime(&end_time);
  std::cout << "Duration: " << duration.count() << "s\n" << std::endl;
}