#include "dual_asym_grouped_gemm.h"

#include <ATen/cuda/CUDAContext.h>
#include <c10/util/BFloat16.h>
#include <c10/cuda/CUDAStream.h>
#include <torch/extension.h>

#include "cutlass/bfloat16.h"
#include "cutlass/complex.h"
#include "cutlass/gemm/kernel/gemm_grouped.h"
#include "cutlass/gemm/kernel/default_gemm_grouped.h"
#include "cutlass/gemm/device/gemm_grouped.h"

#define NUM_STREAM 4

#define CUDA_CALL(code)                               \
    do                                                \
    {                                                 \
        cudaError_t status = code;                    \
        std::string err = cudaGetErrorString(status); \
        TORCH_CHECK(status == cudaSuccess, err);      \
    } while (0)

#define CUBLAS_CALL(code)                                             \
    do                                                                \
    {                                                                 \
        cublasStatus_t status = code;                                 \
        TORCH_CHECK(status == CUBLAS_STATUS_SUCCESS, "CuBLAS Error"); \
    } while (0)

#define GROUPED_GEMM_STRINGIFY_HELPER(x) #x
#define GROUPED_GEMM_STRINGIFY(x) \
GROUPED_GEMM_STRINGIFY_HELPER(x)

template <typename T>
torch::Tensor CopyToDevice(const std::vector<T> &x, const torch::Device &device)
{
    size_t bytes = x.size() * sizeof(T);
    auto options = torch::TensorOptions().dtype(torch::kInt8).device(device);
    torch::Tensor out = torch::empty(bytes, options);

    CUDA_CALL(cudaMemcpyAsync(out.data_ptr(),
                                x.data(), bytes,
                                cudaMemcpyHostToDevice,
                                c10::cuda::getCurrentCUDAStream()));
    return out;
}

using DualExpertGemmKernelNN = typename cutlass::gemm::kernel::DefaultGemmGrouped<
    ::cutlass::bfloat16_t,
    ::cutlass::layout::RowMajor,
    ::cutlass::ComplexTransform::kNone,
    8,
    ::cutlass::bfloat16_t,
    ::cutlass::layout::RowMajor,
    ::cutlass::ComplexTransform::kNone,
    8,
    ::cutlass::bfloat16_t,
    ::cutlass::layout::RowMajor,
    float,
    ::cutlass::arch::OpClassTensorOp,
    ::cutlass::arch::Sm80,
    ::cutlass::gemm::GemmShape<128, 128, 32>,
    ::cutlass::gemm::GemmShape<64, 64, 32>,
    ::cutlass::gemm::GemmShape<16, 8, 16>,
    ::cutlass::epilogue::thread::LinearCombination<::cutlass::bfloat16_t, 8, float, float>,
    ::cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle,
    4>::GemmKernel;

using DualExpertGemmKernelTN = typename cutlass::gemm::kernel::DefaultGemmGrouped<
    ::cutlass::bfloat16_t,
    ::cutlass::layout::ColumnMajor,
    ::cutlass::ComplexTransform::kNone,
    8,
    ::cutlass::bfloat16_t,
    ::cutlass::layout::RowMajor,
    ::cutlass::ComplexTransform::kNone,
    8,
    ::cutlass::bfloat16_t,
    ::cutlass::layout::RowMajor,
    float,
    ::cutlass::arch::OpClassTensorOp,
    ::cutlass::arch::Sm80,
    ::cutlass::gemm::GemmShape<128, 128, 32>,
    ::cutlass::gemm::GemmShape<64, 64, 32>,
    ::cutlass::gemm::GemmShape<16, 8, 16>,
    ::cutlass::epilogue::thread::LinearCombination<::cutlass::bfloat16_t, 8, float, float>,
    ::cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle,
    4>::GemmKernel;

using DualExpertGemmKernelNT = typename cutlass::gemm::kernel::DefaultGemmGrouped<
    ::cutlass::bfloat16_t,
    ::cutlass::layout::RowMajor,
    ::cutlass::ComplexTransform::kNone,
    8,
    ::cutlass::bfloat16_t,
    ::cutlass::layout::ColumnMajor,
    ::cutlass::ComplexTransform::kNone,
    8,
    ::cutlass::bfloat16_t,
    ::cutlass::layout::RowMajor,
    float,
    ::cutlass::arch::OpClassTensorOp,
    ::cutlass::arch::Sm80,
    ::cutlass::gemm::GemmShape<128, 128, 32>,
    ::cutlass::gemm::GemmShape<64, 64, 32>,
    ::cutlass::gemm::GemmShape<16, 8, 16>,
    ::cutlass::epilogue::thread::LinearCombination<::cutlass::bfloat16_t, 8, float, float>,
    ::cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle,
    4>::GemmKernel;

using DualExpertGemmKernelTT = typename cutlass::gemm::kernel::DefaultGemmGrouped<
    ::cutlass::bfloat16_t,
    ::cutlass::layout::ColumnMajor,
    ::cutlass::ComplexTransform::kNone,
    8,
    ::cutlass::bfloat16_t,
    ::cutlass::layout::ColumnMajor,
    ::cutlass::ComplexTransform::kNone,
    8,
    ::cutlass::bfloat16_t,
    ::cutlass::layout::RowMajor,
    float,
    ::cutlass::arch::OpClassTensorOp,
    ::cutlass::arch::Sm80,
    ::cutlass::gemm::GemmShape<128, 128, 32>,
    ::cutlass::gemm::GemmShape<64, 64, 32>,
    ::cutlass::gemm::GemmShape<16, 8, 16>,
    ::cutlass::epilogue::thread::LinearCombination<::cutlass::bfloat16_t, 8, float, float>,
    ::cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle,
    4>::GemmKernel;

using DualExpertGemmNN = ::cutlass::gemm::device::GemmGrouped<DualExpertGemmKernelNN>;
using DualExpertGemmTN = ::cutlass::gemm::device::GemmGrouped<DualExpertGemmKernelTN>;
using DualExpertGemmNT = ::cutlass::gemm::device::GemmGrouped<DualExpertGemmKernelNT>;
using DualExpertGemmTT = ::cutlass::gemm::device::GemmGrouped<DualExpertGemmKernelTT>;

template <typename Gemm>
typename Gemm::Arguments MakeAsymmetricArgumentsSeparated(
    torch::Tensor input_expert0,
    torch::Tensor input_expert1,
    torch::Tensor weight_expert0,
    torch::Tensor weight_expert1,
    torch::Tensor output_expert0,
    torch::Tensor output_expert1)
{

    TORCH_CHECK(input_expert0.dim() == 2 && input_expert1.dim() == 2,
                "Input tensors must be 2D");
    TORCH_CHECK(weight_expert0.dim() == 2 && weight_expert1.dim() == 2,
                "Weight tensors must be 2D");
    TORCH_CHECK(output_expert0.dim() == 2 && output_expert1.dim() == 2,
                "Output tensors must be 2D");

    using LayoutA = typename Gemm::LayoutA;
    using LayoutB = typename Gemm::LayoutB;
    using LayoutC = typename Gemm::LayoutC;

    bool a_is_column_major = std::is_same_v<LayoutA, cutlass::layout::ColumnMajor>;
    bool b_is_column_major = std::is_same_v<LayoutB, cutlass::layout::ColumnMajor>;

    int64_t m0, k0, n0;
    if (a_is_column_major)
    {
        m0 = input_expert0.size(1);
        k0 = input_expert0.size(0);
    }
    else
    {
        m0 = input_expert0.size(0);
        k0 = input_expert0.size(1);
    }

    if (b_is_column_major)
    {
        n0 = weight_expert0.size(0);
        TORCH_CHECK(weight_expert0.size(1) == k0, "Expert 0: k dimensions must match");
    }
    else
    {
        n0 = weight_expert0.size(1);
        TORCH_CHECK(weight_expert0.size(0) == k0, "Expert 0: k dimensions must match");
    }

    int64_t m1, k1, n1;
    if (a_is_column_major)
    {
        m1 = input_expert1.size(1);
        k1 = input_expert1.size(0);
    }
    else
    {
        m1 = input_expert1.size(0);
        k1 = input_expert1.size(1);
    }

    if (b_is_column_major)
    {
        n1 = weight_expert1.size(0);
        TORCH_CHECK(weight_expert1.size(1) == k1, "Expert 1: k dimensions must match");
    }
    else
    {
        n1 = weight_expert1.size(1);
        TORCH_CHECK(weight_expert1.size(0) == k1, "Expert 1: k dimensions must match");
    }

    std::vector<cutlass::gemm::GemmCoord> problem_sizes_host(2);
    problem_sizes_host[0] = cutlass::gemm::GemmCoord(m0, n0, k0);
    problem_sizes_host[1] = cutlass::gemm::GemmCoord(m1, n1, k1);

    int64_t num_experts = 2;
    int threadblock_count = Gemm::sufficient(problem_sizes_host.data(), num_experts);
    if (!threadblock_count)
    {
        TORCH_CHECK(false, "Dual Expert Grouped GEMM execution not possible with HW");
    }

    std::vector<int64_t> lda_host(num_experts);
    std::vector<int64_t> ldb_host(num_experts);
    std::vector<int64_t> ldc_host(num_experts);

    using ElementA = typename Gemm::ElementA;
    using ElementB = typename Gemm::ElementB;
    using ElementC = typename Gemm::ElementC;

    std::vector<ElementA *> ptr_a_host(num_experts);
    std::vector<ElementB *> ptr_b_host(num_experts);
    std::vector<ElementC *> ptr_c_host(num_experts);

    auto problem_0 = problem_sizes_host[0];
    lda_host[0] = LayoutA::packed({problem_0.m(), problem_0.k()}).stride(0);
    ldb_host[0] = LayoutB::packed({problem_0.k(), problem_0.n()}).stride(0);
    ldc_host[0] = LayoutC::packed({problem_0.m(), problem_0.n()}).stride(0);

    ptr_a_host[0] = (ElementA *)input_expert0.data_ptr();
    ptr_b_host[0] = (ElementB *)weight_expert0.data_ptr();
    ptr_c_host[0] = (ElementC *)output_expert0.data_ptr();

    auto problem_1 = problem_sizes_host[1];
    lda_host[1] = LayoutA::packed({problem_1.m(), problem_1.k()}).stride(0);
    ldb_host[1] = LayoutB::packed({problem_1.k(), problem_1.n()}).stride(0);
    ldc_host[1] = LayoutC::packed({problem_1.m(), problem_1.n()}).stride(0);

    ptr_a_host[1] = (ElementA *)input_expert1.data_ptr();
    ptr_b_host[1] = (ElementB *)weight_expert1.data_ptr();
    ptr_c_host[1] = (ElementC *)output_expert1.data_ptr();

    torch::Tensor lda = CopyToDevice(lda_host, input_expert0.device());
    torch::Tensor ldb = CopyToDevice(ldb_host, input_expert0.device());
    torch::Tensor ldc = CopyToDevice(ldc_host, input_expert0.device());
    torch::Tensor ptr_a = CopyToDevice(ptr_a_host, input_expert0.device());
    torch::Tensor ptr_b = CopyToDevice(ptr_b_host, input_expert0.device());
    torch::Tensor ptr_c = CopyToDevice(ptr_c_host, input_expert0.device());
    torch::Tensor problem_sizes = CopyToDevice(problem_sizes_host, input_expert0.device());

    typename Gemm::EpilogueOutputOp::Params epilogue_op(/*alpha=*/1.0f, /*beta=*/0.0f);
    typename Gemm::Arguments arguments(
        (cutlass::gemm::GemmCoord *)problem_sizes.data_ptr(),
        (int)num_experts,
        (int)threadblock_count,
        epilogue_op,
        (ElementA **)ptr_a.data_ptr(),
        (ElementB **)ptr_b.data_ptr(),
        (ElementC **)ptr_c.data_ptr(),
        (ElementC **)ptr_c.data_ptr(),
        (int64_t *)lda.data_ptr(),
        (int64_t *)ldb.data_ptr(),
        (int64_t *)ldc.data_ptr(),
        (int64_t *)ldc.data_ptr(),
        (cutlass::gemm::GemmCoord *)problem_sizes_host.data());

    return arguments;
}

template <typename Gemm>
void executeDualExpertGemm(
    torch::Tensor input_expert0,
    torch::Tensor input_expert1,
    torch::Tensor weight_expert0,
    torch::Tensor weight_expert1,
    torch::Tensor output_expert0,
    torch::Tensor output_expert1)
{

    Gemm gemm;

    auto arguments = MakeAsymmetricArgumentsSeparated<Gemm>(
        input_expert0, input_expert1,
        weight_expert0, weight_expert1,
        output_expert0, output_expert1);

    int64_t workspace_size = gemm.get_workspace_size(arguments);
    auto options = torch::TensorOptions().dtype(torch::kInt8).device(input_expert0.device());
    torch::Tensor workspace = torch::empty(workspace_size, options);

    if (gemm.initialize(arguments, workspace.data_ptr()) != cutlass::Status::kSuccess)
    {
        TORCH_CHECK(false, "Failed to initialize CUTLASS Asymmetric Dual Expert GEMM");
    }

    if (gemm.run(c10::cuda::getCurrentCUDAStream()) != cutlass::Status::kSuccess)
    {
        TORCH_CHECK(false, "Failed to run CUTLASS Asymmetric Dual Expert GEMM");
    }
}

void AsymmetricDualExpertGemm(
    torch::Tensor input_expert0,
    torch::Tensor input_expert1,
    torch::Tensor weight_expert0,
    torch::Tensor weight_expert1,
    torch::Tensor output_expert0,
    torch::Tensor output_expert1,
    bool trans_a, bool trans_b)
{

    TORCH_CHECK(input_expert0.device() == input_expert1.device() &&
                    input_expert0.device() == weight_expert0.device() &&
                    input_expert0.device() == weight_expert1.device() &&
                    input_expert0.device() == output_expert0.device() &&
                    input_expert0.device() == output_expert1.device(),
                "All tensors must be on the same device");

    TORCH_CHECK(input_expert0.device().is_cuda(),
                "All tensors must be on CUDA device for CUTLASS GEMM");

    TORCH_CHECK(input_expert0.dtype() == torch::kBFloat16,
                "All tensors must be BFloat16 for this kernel");

    torch::Tensor input0_contiguous = input_expert0.contiguous();
    torch::Tensor input1_contiguous = input_expert1.contiguous();
    torch::Tensor weight0_contiguous = weight_expert0.contiguous();
    torch::Tensor weight1_contiguous = weight_expert1.contiguous();
    torch::Tensor output0_contiguous = output_expert0.contiguous();
    torch::Tensor output1_contiguous = output_expert1.contiguous();

    if (!trans_a && !trans_b)
    {
        using Gemm = DualExpertGemmNN;
        executeDualExpertGemm<Gemm>(input0_contiguous, input1_contiguous,
                                    weight0_contiguous, weight1_contiguous,
                                    output0_contiguous, output1_contiguous);
    }
    else if (trans_a && !trans_b)
    {
        using Gemm = DualExpertGemmTN;
        executeDualExpertGemm<Gemm>(input0_contiguous, input1_contiguous,
                                    weight0_contiguous, weight1_contiguous,
                                    output0_contiguous, output1_contiguous);
    }
    else if (!trans_a && trans_b)
    {
        using Gemm = DualExpertGemmNT;
        executeDualExpertGemm<Gemm>(input0_contiguous, input1_contiguous,
                                    weight0_contiguous, weight1_contiguous,
                                    output0_contiguous, output1_contiguous);
    }
    else
    {
        using Gemm = DualExpertGemmTT;
        executeDualExpertGemm<Gemm>(input0_contiguous, input1_contiguous,
                                    weight0_contiguous, weight1_contiguous,
                                    output0_contiguous, output1_contiguous);
    }
}