#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 <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_;

  map<string, vector<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 pre_allocate_stats(size_t num = 100000) {
    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));
    }
  }

  static void read_args(int argc, char **argv, int &num_threads, string &path) {
    if (argc < 3) {
      cout << "Must Specifiy concurrency and output directory! (usage: `benchmark <output_directory> <num_threads>`)"
           << endl;
      exit(1);
    }

    string tmp = argv[1];
    path = tmp;
    num_threads = atoi(argv[2]);
  }

  void start_iteration() {
    if (memory_stats_enabled_) {
      auto memory_stats = query_process_memory_pages();
      memory_pre_run_ = memory_stats.first;
    }

    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_;
    }
  }

  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,
                      int warmup_count,
                      const function<void(void)> prepare = []() {},
                      const function<void(void)> finish = []() {}) {
    for (int i = 0; i < warmup_count; i++) {
      prepare();
      measure();
    }

    for (int i = 0; i < count; i++) {
      using namespace std::literals;
      this_thread::sleep_for(100us);
      prepare();
      start_iteration();
      measure();
      end_iteration();
      finish();
    }
  }

  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