#include "maxwell_lib.hpp"
#include <chrono>
#include <iostream>
#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

#ifdef MAXWELL_CUDA
extern "C" void launch_maxwell_cuda(const float *h_A, const uint8_t *h_B_packed,
                                    float *h_C, const BlockCoord *h_active,
                                    int num_active, int M, int K, int N);
extern "C" void launch_maxwell_cuda_wmma(const float *h_A,
                                         const uint8_t *h_B_packed, float *h_C,
                                         const BlockCoord *h_active,
                                         int num_active, int M, int K, int N);
extern "C" void launch_maxwell_backward_weights(const float *h_A, const float *h_grad_Y,
                                                float *h_grad_W, int M, int K, int N);
#endif

namespace py = pybind11;

PackedWeights create_packed_weights(py::array_t<int8_t> input) {
  py::buffer_info buf = input.request();
  int rows = buf.shape[0];
  int cols = buf.shape[1];
  PackedWeights pw(rows, cols);
  auto ptr = static_cast<int8_t *>(buf.ptr);
  for (int i = 0; i < rows; ++i) {
    for (int j = 0; j < cols; ++j) {
      pw.set(i, j, ptr[i * cols + j]);
    }
  }
  return pw;
}

double run_maxwell_cuda(py::array_t<float> A, const PackedWeights &B,
                        py::array_t<float> C) {
#ifdef MAXWELL_CUDA
  auto bufA = A.request();
  auto bufC = C.request();
  float *ptrA = static_cast<float *>(bufA.ptr);
  float *ptrC = static_cast<float *>(bufC.ptr);
  int M = bufA.shape[0];
  int K = bufA.shape[1];
  int N = bufC.shape[1];

  auto start = std::chrono::high_resolution_clock::now();
  std::vector<BlockCoord> active_blocks = invoke_demon_lib(ptrA, M, K);
  launch_maxwell_cuda(ptrA, B.data.data(), ptrC, active_blocks.data(),
                      active_blocks.size(), M, K, N);
  auto end = std::chrono::high_resolution_clock::now();
  return std::chrono::duration<double>(end - start).count();
#else
  throw std::runtime_error("CUDA not compiled");
#endif
}

void run_maxwell_backward_weights_cuda(py::array_t<float> A, py::array_t<float> grad_Y,
                                       py::array_t<float> grad_W) {
#ifdef MAXWELL_CUDA
    auto bufA = A.request();
    auto bufGY = grad_Y.request();
    auto bufGW = grad_W.request();
    
    int M = bufA.shape[0];
    int K = bufA.shape[1];
    int N = bufGY.shape[1];
    
    launch_maxwell_backward_weights(
        static_cast<float*>(bufA.ptr),
        static_cast<float*>(bufGY.ptr),
        static_cast<float*>(bufGW.ptr),
        M, K, N
    );
#else
    throw std::runtime_error("CUDA not compiled");
#endif
}

PYBIND11_MODULE(maxwell_core, m) {
  m.doc() = "Maxwell Core C++ SOTA Backend";
  
  py::class_<PackedWeights>(m, "PackedWeights")
      .def(py::init<int, int>());

  m.def("create_packed_weights", &create_packed_weights);
  m.def("run_maxwell_cuda", &run_maxwell_cuda);
  m.def("run_maxwell_backward_weights_cuda", &run_maxwell_backward_weights_cuda);
}
