#include <pls/pls.h>
#include <pls/internal/helpers/profiler.h>
#include <pls/internal/helpers/mini_benchmark.h>

const int MATRIX_SIZE = 128;

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

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

  void multiply(const matrix<T, SIZE> &a, const matrix<T, SIZE> &b) {
    pls::algorithm::parallel_for(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;
    }
  }
}

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

  pls::internal::helpers::run_mini_benchmark([&] {
    result.multiply(a, b);
  }, 8, 1000);

  PROFILE_SAVE("test_profile.prof")
}

//int main() {
//  PROFILE_ENABLE
//  pls::malloc_scheduler_memory my_scheduler_memory{8, 2u << 18};
//  pls::scheduler scheduler{&my_scheduler_memory, 8};
//
//  matrix<double, MATRIX_SIZE> a;
//  matrix<double, MATRIX_SIZE> b;
//  matrix<double, MATRIX_SIZE> result;
//  fill_with_data(a, b);
//
//  scheduler.perform_work([&] {
//    PROFILE_MAIN_THREAD
//    for (int i = 0; i < 5000; i++) {
//      PROFILE_WORK_BLOCK("Top Level")
//      result.multiply(a, b);
//    }
//  });
//
//  PROFILE_SAVE("test_profile.prof")
//}
//