#ifndef BENCHMARK_RUNNER_H
#define BENCHMARK_RUNNER_H

#include <string>
#include <cstdlib>
#include <vector>
#include <chrono>
#include <numeric>
#include <iostream>
#include <fstream>
#include <bits/stdc++.h>
#include <thread>
#include <map>
#include <time.h>

#include <tuple>
#include <unistd.h>
#include <string>
#include <fstream>
#include <sstream>

using namespace std;

class benchmark_runner {
 private:
  string csv_path_;
  string csv_name_;

  chrono::steady_clock::time_point last_start_time_;
  vector<long> times_;

  bool memory_stats_enabled_{false};
  const string MEMORY_PRE_RUN = "memory_pages_pre_run";
  const string MEMORY_POST_RUN = "memory_pages_post_run";
  unsigned long memory_pre_run_;
  unsigned long memory_post_run_;

  bool wall_time_enabled_{false};
  const string WALL_TIME_PRE_RUN = "wall_time_pre_run_us";
  const string WALL_TIME_POST_RUN = "wall_time_post_run_us";
  unsigned long wall_time_pre_run_;
  unsigned long wall_time_post_run_;

  const string WALL_TIME_ITERATION_START = "wall_time_iteration_start_us";
  const string WALL_TIME_ITERATION_END = "wall_time_iteration_end_us";

  map<string, vector<unsigned long>> custom_stats_;

  void print_statistics() {
    long time_sum = std::accumulate(times_.begin(), times_.end(), 0l);
    cout << "Average Runtime (us): " << (time_sum / times_.size()) << endl;

    for (auto &iter : custom_stats_) {
      long custom_stat_sum = std::accumulate(iter.second.begin(), iter.second.end(), 0l);
      cout << "Average " << iter.first << ": " << (custom_stat_sum / times_.size()) << endl;
    }
  }

  inline bool file_exists(const std::string &name) {
    ifstream f(name);
    return f.good();
  }

 public:
  benchmark_runner(string csv_path, string csv_name) : csv_path_{std::move(csv_path)},
                                                       csv_name_{std::move(csv_name)},
                                                       times_{} {
    string command = "mkdir -p " + csv_path_;
    int res = system(command.c_str());
    if (res) {
      cout << "Error while creating directory!" << endl;
      exit(1);
    }
  }

  /**
   * Queries the pages used by the current process.
   * Splits between private and shared pages.
   *
   * @return Tuple(private_pages, shared_pages) used by the process.
   */
  static std::pair<unsigned long, unsigned long> query_process_memory_pages() {
    pid_t my_pid = getpid();
    std::ostringstream proc_stats_path_builder;
    proc_stats_path_builder << "/proc/" << my_pid << "/smaps_rollup";
    std::string proc_stats_path = proc_stats_path_builder.str();

    std::ifstream proc_stats_file{proc_stats_path};

    std::string line;
    unsigned long total_shared = 0;
    unsigned long total_private = 0;
    while (std::getline(proc_stats_file, line)) {
      std::stringstream line_input(line);
      std::string type;
      unsigned long size;
      line_input >> type >> size;

      if (type.find("Shared") != std::string::npos) {
        total_shared += size;
      }
      if (type.find("Private") != std::string::npos) {
        total_private += size;
      }
    }

    return std::make_pair(total_private, total_shared);
  }

  void add_custom_stats_field(const string name) {
    custom_stats_.insert({name, {}});
  }

  void enable_memory_stats() {
    memory_stats_enabled_ = true;
    add_custom_stats_field(MEMORY_PRE_RUN);
    add_custom_stats_field(MEMORY_POST_RUN);
  }

  void enable_wall_time_stats() {
    wall_time_enabled_ = true;
    add_custom_stats_field(WALL_TIME_PRE_RUN);
    add_custom_stats_field(WALL_TIME_POST_RUN);
    add_custom_stats_field(WALL_TIME_ITERATION_START);
    add_custom_stats_field(WALL_TIME_ITERATION_END);
  }

  void pre_allocate_stats(size_t num = 500000) {
    times_.reserve(num);
    memset(times_.data(), 'a', num * sizeof(long));
    for (auto &iter : custom_stats_) {
      iter.second.reserve(num);
      memset(iter.second.data(), 'a', num * sizeof(long));
    }
  }

  struct benchmark_settings {
    enum TYPE { ISOLATED, PERIODIC };

    string output_directory_;
    size_t size_;
    unsigned num_threads_;

    TYPE type_;

    size_t iterations_;

    unsigned long interval_period_;
    unsigned long interval_deadline_;
  };

  static benchmark_settings parse_parameters(int argc, char **argv) {
    benchmark_settings result;
    string tmp;

    if (argc != 5 && argc != 7) {
      printf("usage 1: `benchmark <output_directory> <size> <num_threads> <iterations>`\n");
      printf("usage 2: `benchmark <output_directory> <size> <num_threads> <iterations> <period> <deadline>`\n");
      exit(1);
    }

    result.output_directory_ = argv[1];
    tmp = argv[2];
    result.size_ = std::stoi(tmp);
    tmp = argv[3];
    result.num_threads_ = std::stoi(tmp);

    if (argc == 5) {
      result.type_ = benchmark_settings::ISOLATED;

      tmp = argv[4];
      result.iterations_ = std::stoi(tmp);
    } else {
      result.type_ = benchmark_settings::PERIODIC;

      tmp = argv[4];
      result.iterations_ = std::stoi(tmp);
      tmp = argv[5];
      result.interval_period_ = std::stoi(tmp);
      tmp = argv[6];
      result.interval_deadline_ = std::stoi(tmp);
    }

    return result;
  }

  void start_iteration() {
    if (memory_stats_enabled_) {
      auto memory_stats = query_process_memory_pages();
      memory_pre_run_ = memory_stats.first;
    }
    if (wall_time_enabled_) {
      wall_time_pre_run_ =
          chrono::duration_cast<chrono::microseconds>(chrono::system_clock::now().time_since_epoch()).count();
    }

    last_start_time_ = chrono::steady_clock::now();
  }

  void end_iteration() {
    size_t iteration_index = times_.size();

    auto end_time = chrono::steady_clock::now();
    long time = chrono::duration_cast<chrono::microseconds>(end_time - last_start_time_).count();

    times_.emplace_back(time);
    for (auto &iter : custom_stats_) {
      iter.second.emplace_back(0);
    }
    if (memory_stats_enabled_) {
      auto memory_stats = query_process_memory_pages();
      memory_post_run_ = memory_stats.first;

      custom_stats_[MEMORY_PRE_RUN][iteration_index] = memory_pre_run_;
      custom_stats_[MEMORY_POST_RUN][iteration_index] = memory_post_run_;
    }
    if (wall_time_enabled_) {
      wall_time_post_run_ =
          chrono::duration_cast<chrono::microseconds>(chrono::system_clock::now().time_since_epoch()).count();

      custom_stats_[WALL_TIME_PRE_RUN][iteration_index] = wall_time_pre_run_;
      custom_stats_[WALL_TIME_POST_RUN][iteration_index] = wall_time_post_run_;
    }
  }

  void store_custom_stat(const string &name, long value) {
    auto &stat_vector = custom_stats_[name];
    stat_vector[stat_vector.size() - 1] = value;
  }

  void run_iterations(int count,
                      const function<void(void)> measure,
                      const function<void(void)> prepare = []() {},
                      const function<void(void)> finish = []() {}) {
    for (int i = 0; i < count; i++) {
      using namespace std::literals;
      this_thread::sleep_for(100us);
      prepare();
      start_iteration();
      measure();
      end_iteration();
      finish();
    }
  }

  void add_to_timespec(timespec &timespec, size_t seconds, size_t nanoseconds) {
    timespec.tv_sec += seconds;
    timespec.tv_nsec += nanoseconds;

    while (timespec.tv_nsec > 1000l * 1000l * 1000l) {
      timespec.tv_nsec -= 1000l * 1000l * 1000l;
      timespec.tv_sec++;
    }
  }

  void run_periodic(size_t count,
                    size_t period_us,
                    size_t deadline_us,
                    const function<void(void)> measure) {
    size_t period_nanoseconds = 1000lu * period_us;
    size_t period_seconds = period_nanoseconds / (1000lu * 1000lu * 1000lu);
    period_nanoseconds = period_nanoseconds - period_seconds * 1000lu * 1000lu * 1000lu;

    size_t deadline_nanoseconds = 1000lu * deadline_us;
    size_t deadline_seconds = deadline_nanoseconds / (1000lu * 1000lu * 1000lu);
    deadline_nanoseconds = deadline_nanoseconds - deadline_seconds * 1000lu * 1000lu * 1000lu;

    // Prepare basic time spec for first iteration
    timespec iteration_start, iteration_end, deadline_end, finish_time;
    if (clock_gettime(CLOCK_MONOTONIC, &iteration_start) == -1) {
      perror("clock_gettime");
      exit(1);
    }
    add_to_timespec(iteration_start, period_seconds, period_nanoseconds);

    size_t current_iteration = 0;
    size_t deadline_misses = 0;
    while (current_iteration < count) {
      // Sleep until the next iteration
      long sleep_error = clock_nanosleep(CLOCK_MONOTONIC, TIMER_ABSTIME, &iteration_start, nullptr);
      if (sleep_error) {
        printf("Sleep Error %ld\n", sleep_error);
      }

      // Invoke iteration
      start_iteration();
      measure();
      end_iteration();

      // Calculate all relevant time points for this iteration
      if (clock_gettime(CLOCK_MONOTONIC, &finish_time) == -1) {
        perror("clock_gettime");
        exit(1);
      }

      iteration_end = iteration_start;
      add_to_timespec(iteration_end, period_seconds, period_nanoseconds);

      deadline_end = iteration_start;
      add_to_timespec(deadline_end, deadline_seconds, deadline_nanoseconds);

      // Keep stats...
      if (wall_time_enabled_) {
        custom_stats_[WALL_TIME_ITERATION_START][current_iteration] =
            iteration_start.tv_sec * 1000 * 1000 + iteration_start.tv_nsec / 1000;
        custom_stats_[WALL_TIME_ITERATION_END][current_iteration] =
            iteration_end.tv_sec * 1000 * 1000 + iteration_end.tv_nsec / 1000;
      }

      // Store 'actual' wall time instead of iteration time (we want to include sleeping here!)
      unsigned long wall_time_us = 0;
      wall_time_us += (finish_time.tv_sec - iteration_start.tv_sec) * 1000l * 1000l;
      long nano_second_difference = ((long) finish_time.tv_nsec - (long) iteration_start.tv_nsec) / 1000l;
      wall_time_us += nano_second_difference;
      times_[current_iteration] = wall_time_us;

      if (finish_time.tv_sec >= deadline_end.tv_sec && finish_time.tv_nsec > deadline_end.tv_nsec) {
        deadline_misses++;
      }

      // Skip iterations if their start time is later than the current time (skipping)
      while (finish_time.tv_sec > iteration_end.tv_sec
          || (finish_time.tv_sec == iteration_end.tv_sec && finish_time.tv_nsec > iteration_end.tv_nsec)) {
        iteration_start = iteration_end;

        iteration_end = iteration_start;
        add_to_timespec(iteration_end, period_seconds, period_nanoseconds);

        current_iteration++;
        start_iteration();
        end_iteration();
        times_[current_iteration] = 0;
      }
      // Progress to next iteration (normally)
      current_iteration++;
      iteration_start = iteration_end;
    }

    printf("%ld deadline misses!\n", deadline_misses);
  }

  void commit_results(bool print_stats) {
    if (print_stats) {
      print_statistics();
    }

    string full_filename = csv_path_ + csv_name_;
    bool write_header = !file_exists(full_filename);

    { // Scope for output file
      ofstream o(full_filename, std::fstream::out | std::fstream::app);

      if (write_header) {
        o << "runtime_us";
        for (auto &iter : custom_stats_) {
          o << ";" << iter.first;
        }
        o << endl;
      }

      // TODO: make this more efficient
      for (size_t i = 0; i < times_.size(); i++) {
        o << times_[i];
        for (auto &iter : custom_stats_) {
          o << ";" << iter.second[i];
        }
        o << endl;
      }
    } // End Scope for output file

    times_.clear();
    for (auto &iter : custom_stats_) {
      iter.second.clear();
    }
  }
};

#endif //BENCHMARK_RUNNER_H