
#ifndef PLS_MINI_BENCHMARK_H
#define PLS_MINI_BENCHMARK_H

#include "pls/internal/scheduling/scheduler_memory.h"
#include "pls/internal/scheduling/scheduler.h"

#include <chrono>
#include <iostream>

namespace pls {
namespace internal {
namespace helpers {

// TODO: Clean up (separate into small functions and .cpp file)
template<typename Function>
void run_mini_benchmark(const Function &lambda,
                        size_t max_threads,
                        unsigned long max_runtime_ms = 1000,
                        unsigned long warmup_time_ms = 100) {
  using namespace std;
  using namespace pls::internal::scheduling;

  malloc_scheduler_memory scheduler_memory{max_threads, 2u << 14};
  for (unsigned int num_threads = 1; num_threads <= max_threads; num_threads++) {
    scheduler local_scheduler{&scheduler_memory, num_threads};

    chrono::high_resolution_clock::time_point start_time;
    chrono::high_resolution_clock::time_point end_time;
    long max_local_time = 0;
    long total_time = 0;
    long iterations = 0;

    local_scheduler.perform_work([&] {
      start_time = chrono::high_resolution_clock::now();
      end_time = start_time;
      chrono::high_resolution_clock::time_point planned_end_time = start_time + chrono::milliseconds(max_runtime_ms);
      chrono::high_resolution_clock::time_point planned_warmup_time = start_time + chrono::milliseconds(warmup_time_ms);

      while (end_time < planned_end_time) {
        if (end_time < planned_warmup_time) {
          lambda();
        } else {
          auto local_start_time = chrono::high_resolution_clock::now();
          lambda();
          auto local_end_time = chrono::high_resolution_clock::now();
          long local_time = chrono::duration_cast<chrono::microseconds>(local_end_time - local_start_time).count();

          total_time += local_time;
          max_local_time = std::max(local_time, max_local_time);
          iterations++;
        }
        end_time = chrono::high_resolution_clock::now();
      }
    });
    double time_per_iteration = (double) total_time / iterations;

    std::cout << (long) time_per_iteration << " (" << max_local_time << ")";
    if (num_threads < max_threads) {
      std::cout << "\t\t";
    }
  }
  std::cout << std::endl;
}

}
}
}

#endif //PLS_MINI_BENCHMARK_H
