#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <cstdlib>
#include <iostream>

#include <CL/sycl.hpp>
#include <utility>

#include "matrix.hpp"
#include "timer.hpp"

namespace sycl = cl::sycl;

#define stream_hex(_val) std::hex << _val << std::dec

#define print_pair(_desc, _chksm, _time)                                       \
  std::cout << _desc << "\n\t"                                                 \
            << "\t-> Check: 0x" << stream_hex(_chksm)                          \
            << "\tRuntime: " << _time << " us\n\n"

using data_type = int;

auto matrixMultCPU(const Matrix<data_type> &matA,
                   const Matrix<data_type> &matB) {
  Matrix<data_type> res(matA.rows, matB.cols);
  for (std::uint32_t i = 0; i < res.rows; i++) {
    for (std::uint32_t j = 0; j < res.cols; j++) {
      auto &res_val = res(i, j) = 0;
      for (std::uint32_t k = 0; k < matA.cols; k++) {
        res_val += matA(i, k) * matB(k, j);
      }
    }
  }

  return res.chksum();
}

auto matrixMultTransposeCPU(const Matrix<data_type> &matA,
                            const Matrix<data_type> &matB) {
  Matrix<data_type> matB_t = matB.t();
  Matrix<data_type> res(matA.rows, matB.cols);
  for (std::uint32_t i = 0; i < res.rows; i++) {
    for (std::uint32_t j = 0; j < res.cols; j++) {
      auto &res_val = res(i, j) = 0;
      for (std::uint32_t k = 0; k < matA.cols; k++) {
        res_val += matA(i, k) * matB_t(j, k);
      }
    }
  }

  return res.chksum();
}

auto matrixMultSYCL(sycl::queue &q, const Matrix<data_type> &matA,
                    const Matrix<data_type> &matB) {
  Matrix<data_type> matRes(matA.rows, matB.cols);
  sycl::range<2> global_range(matRes.rows, matRes.cols);

  {
    sycl::buffer<data_type, 2> b_matA(matA.mem.data(),
                                      sycl::range<2>(matA.rows, matA.cols));

    sycl::buffer<data_type, 2> b_matB(matB.mem.data(),
                                      sycl::range<2>(matB.rows, matB.cols));

    sycl::buffer<data_type, 2> b_matRes(
        matRes.mem.data(), sycl::range<2>(matRes.rows, matRes.cols));

    q.submit([&](sycl::handler &h) {
      auto acc_matA = b_matA.get_access<sycl::access::mode::read>(h);
      auto acc_matB = b_matB.get_access<sycl::access::mode::read>(h);
      auto acc_matRes = b_matRes.get_access<sycl::access::mode::write>(h);

      h.parallel_for(global_range, [=](sycl::id<2> ID) {
        auto i = ID[0];
        auto j = ID[1];
        data_type sum = 0;

        for (auto k = 0; k < matA.cols; k++) {
          sum += acc_matA[i][k] * acc_matB[k][j];
        }
        acc_matRes[i][j] = sum;
      });
    });
  }

  q.wait();

  return matRes.chksum();
}

auto matrixMultTransposeSYCL(sycl::queue &q, const Matrix<data_type> &matA,
                             const Matrix<data_type> &matB) {

  Matrix<data_type> matB_t = matB.t();
  Matrix<data_type> matRes(matA.rows, matB.cols);
  sycl::range<2> global_range(matRes.rows, matRes.cols);

  {
    sycl::buffer<data_type, 2> b_matA(matA.mem.data(),
                                      sycl::range<2>(matA.rows, matA.cols));

    sycl::buffer<data_type, 2> b_matB(matB_t.mem.data(),
                                      sycl::range<2>(matB_t.rows, matB_t.cols));

    sycl::buffer<data_type, 2> b_matRes(
        matRes.mem.data(), sycl::range<2>(matRes.rows, matRes.cols));

    q.submit([&](sycl::handler &h) {
      auto acc_matA = b_matA.get_access<sycl::access::mode::read>(h);
      auto acc_matB = b_matB.get_access<sycl::access::mode::read>(h);
      auto acc_matRes = b_matRes.get_access<sycl::access::mode::write>(h);

      h.parallel_for(global_range, [=](sycl::id<2> ID) {
        auto i = ID[0];
        auto j = ID[1];
        data_type sum = 0;

        for (auto k = 0; k < matA.cols; k++) {
          sum += acc_matA[i][k] * acc_matB[j][k];
        }
        acc_matRes[i][j] = sum;
      });
    });
  }

  q.wait();

  return matRes.chksum();
}

/* Obtains the previous power of two from the given integer.
 * It works by masking out all ones after the first one bit,
 * then leaves the first one bit intact, effectively
 * yielding the first power of two < x. */
inline int prevPowerOfTwo(int x) {
  if (x < 0) {
    return 0;
  }
  --x;
  x |= x >> 1;
  x |= x >> 2;
  x |= x >> 4;
  x |= x >> 8;
  x |= x >> 16;
  return x - (x >> 1);
}

template <class T>
auto matrixMultTiledSYCL(sycl::queue &q, const Matrix<T> &matA,
                         const Matrix<T> &matB) {
  Matrix<T> matRes(matA.rows, matB.cols);

  std::size_t max_group_size =
      q.get_device().get_info<sycl::info::device::max_work_group_size>();

  // lets assume we always have a maximum group size with a power of 2
  const std::uint32_t max_block_size =
      static_cast<std::uint32_t>(prevPowerOfTwo(std::sqrt(max_group_size)));

  const std::uint32_t block_size = std::min(matA.cols, max_block_size);

  sycl::range<2> global_range(matRes.rows, matRes.cols);
  sycl::range<2> local_range(block_size, block_size);

  {
    sycl::buffer<T, 2> b_matA(matA.mem.data(),
                              sycl::range<2>(matA.rows, matA.cols));

    sycl::buffer<T, 2> b_matB(matB.mem.data(),
                              sycl::range<2>(matB.rows, matB.cols));

    sycl::buffer<T, 2> b_matRes(matRes.mem.data(),
                                sycl::range<2>(matRes.rows, matRes.cols));

    q.submit([&](sycl::handler &h) {
      auto acc_matA = b_matA.template get_access<sycl::access::mode::read>(h);
      auto acc_matB = b_matB.template get_access<sycl::access::mode::read>(h);
      auto acc_matRes =
          b_matRes.template get_access<sycl::access::mode::write>(h);

      sycl::accessor<int, 2, sycl::access::mode::read_write,
                     sycl::access::target::local>
          tileA(local_range, h);

      sycl::accessor<int, 2, sycl::access::mode::read_write,
                     sycl::access::target::local>
          tileB(local_range, h);

      h.parallel_for<class tiled_matmul>(
          sycl::nd_range{global_range, local_range}, [=](sycl::nd_item<2> &ID) {
            const int i = ID.get_global_id(0);
            const int j = ID.get_global_id(1);

            const int local_i = ID.get_local_id(0);
            const int local_j = ID.get_local_id(1);

            const int max_tile = ID.get_group_range(0);

            // Current local IDem
            T sum = 0;

            for (int tile_i = 0; tile_i < max_tile; tile_i++) {
              tileA[local_i][local_j] =
                  acc_matA[i][tile_i * block_size + local_j];
              tileB[local_j][local_i] =
                  acc_matB[block_size * tile_i + local_i][j];

              ID.barrier(sycl::access::fence_space::local_space);

              for (auto k = 0; k < block_size; k++) {
                sum += tileA[local_i][k] * tileB[local_j][k];
              }

              ID.barrier(sycl::access::fence_space::local_space);
            }

            acc_matRes[i][j] = sum;
          });
    });
  }

  q.wait();

  return matRes.chksum();
}

auto main(int argc, char **argv) -> int {

  if (argc != 3) {
    std::cerr << "Provide 2 arguments to the program!\n"
              << "Usage: <prog> <matA>.txt <matB>.txt\n";
    return EXIT_FAILURE;
  }

  Matrix<data_type> matA(argv[1]);
  Matrix<data_type> matB(argv[2]);

  assert(matA.rows == matB.cols);

#ifdef SEQ_BENCH
  auto cpu_chksum = measure<>::duration(matrixMultCPU, matA, matB);
  print_pair("CPU - naive", cpu_chksum.first, cpu_chksum.second.count());

  auto cpu_transp_chksum =
      measure<>::duration(matrixMultTransposeCPU, matA, matB);
  print_pair("CPU - transposed", cpu_transp_chksum.first,
             cpu_transp_chksum.second.count());
#endif

  sycl::queue cpu_queue(sycl::cpu_selector_v);

  auto omp_chksum = measure<>::duration(matrixMultSYCL, cpu_queue, matA, matB);
  print_pair("OMP - naive", omp_chksum.first, omp_chksum.second.count());

  auto omp_transp_chksum =
      measure<>::duration(matrixMultTransposeSYCL, cpu_queue, matA, matB);
  print_pair("OMP - transposed", omp_transp_chksum.first,
             omp_transp_chksum.second.count());

  sycl::queue gpu_queue(sycl::gpu_selector_v);

  auto gpu_chksum = measure<>::duration(matrixMultSYCL, gpu_queue, matA, matB);
  print_pair("GPU - naive", gpu_chksum.first, gpu_chksum.second.count());

  auto gpu_transp_chksum =
      measure<>::duration(matrixMultTransposeSYCL, gpu_queue, matA, matB);
  print_pair("GPU - transposed", gpu_transp_chksum.first,
             gpu_transp_chksum.second.count());

  auto gpu_tiled_chksum = measure<>::duration(matrixMultTiledSYCL<data_type>,
                                              gpu_queue, matA, matB);
  print_pair("GPU - tiled", gpu_tiled_chksum.first,
             gpu_tiled_chksum.second.count());

  return EXIT_SUCCESS;
}
