#include "pls/internal/scheduling/scheduler.h"
#include "pls/internal/scheduling/parallel_result.h"
#include "pls/internal/scheduling/scheduler_memory.h"
#include "pls/algorithms/for_each.h"
using namespace pls::internal::scheduling;

#include <chrono>

const int MATRIX_SIZE = 128;

template<typename T, int SIZE>
class matrix {
 public:
  T data[SIZE][SIZE];

  explicit matrix(T i = 1) {
    std::fill(&data[0][0], &data[0][0] + SIZE * SIZE, i);
  }

  parallel_result<int> multiply(const matrix<T, SIZE> &a, const matrix<T, SIZE> &b) {
    return pls::algorithm::for_each_range(0, SIZE, [&](int i) {
      this->multiply_column(i, a, b);
    });
  }

 private:
  void multiply_column(int i, const matrix<T, SIZE> &a, const matrix<T, SIZE> &b) {
    for (int j = 0; j < SIZE; ++j) {
      data[i][j] = 0;
    }
    for (int k = 0; k < SIZE; ++k) {
      for (int j = 0; j < SIZE; ++j) {
        data[i][j] += a.data[i][k] * b.data[k][j];
      }
    }
  }
};

void fill_with_data(matrix<double, MATRIX_SIZE> &a, matrix<double, MATRIX_SIZE> &b) {
  // Fill in some data...
  for (int i = 0; i < MATRIX_SIZE; i++) {
    for (int j = 0; j < MATRIX_SIZE; j++) {
      a.data[i][j] = i;
      b.data[i][j] = j;
    }
  }
}

static constexpr int NUM_ITERATIONS = 1000;
constexpr size_t NUM_THREADS = 3;

constexpr size_t NUM_TASKS = 128;

constexpr size_t NUM_CONTS = 128;
constexpr size_t MAX_CONT_SIZE = 512;

int main() {
  PROFILE_ENABLE
  matrix<double, MATRIX_SIZE> a;
  matrix<double, MATRIX_SIZE> b;
  matrix<double, MATRIX_SIZE> result;
  fill_with_data(a, b);

  static_scheduler_memory<NUM_THREADS,
                          NUM_TASKS,
                          NUM_CONTS,
                          MAX_CONT_SIZE> static_scheduler_memory;

  scheduler scheduler{static_scheduler_memory, NUM_THREADS};

  auto start = std::chrono::steady_clock::now();
  for (int i = 0; i < NUM_ITERATIONS; i++) {
    scheduler.perform_work([&]() {
      PROFILE_MAIN_THREAD;
      return scheduler::par([&]() {
        return result.multiply(a, b);
      }, []() {
        return parallel_result<int>{0};
      }).then([](int, int) {
        return parallel_result<int>{0};
      });
    });
  }
  auto end = std::chrono::steady_clock::now();
  std::cout << "Framework:  " << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count()
            << std::endl;
}

//int main() {
//  PROFILE_ENABLE
//  pls::malloc_scheduler_memory my_scheduler_memory{8, 2u << 18u};
//  pls::scheduler scheduler{&my_scheduler_memory, 4};
//
//  matrix<double, MATRIX_SIZE> a;
//  matrix<double, MATRIX_SIZE> b;
//  matrix<double, MATRIX_SIZE> result;
//  fill_with_data(a, b);
//
//  scheduler.perform_work([&] {
//    auto start_time = std::chrono::high_resolution_clock::now();
//    PROFILE_MAIN_THREAD
//    for (int i = 0; i < 10000; i++) {
//      PROFILE_WORK_BLOCK("Top Level")
//      result.multiply(a, b);
//    }
//    auto end_time = std::chrono::high_resolution_clock::now();
//    long time = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time).count();
//    std::cout << "Runtime: " << time << "us" << std::endl;
//  });
//
//  PROFILE_SAVE("test_profile.prof")
//}