#include <torch/library.h>
#include <c10/cuda/CUDAStream.h>
#include <c10/cuda/CUDAException.h>
#include <algorithm>

#include "torch_binding.h"

// Not CHECK: torch/library.h pulls in glog's CHECK and redefining it warns.
#define M3_F32(x) TORCH_CHECK((x).is_cuda() && (x).scalar_type() == at::kFloat && (x).is_contiguous(), \
                              #x " must be a contiguous float32 CUDA tensor")

#define STEP_ARGS \
    const void* q, const void* k, const void* v, const void* z, const float* adt, \
    const float* dt, const float* trap, const void* q_bias, const void* k_bias, \
    const float* angles, const void* mimo_v, const void* mimo_o, const void* mimo_z, \
    const void* Dvec, const float* angle_in, float* angle_out, float* S_st, \
    const float* kprev_in, const float* vprev_in, float* kprev_out, float* vprev_out, \
    void* y, int B, int H, int Gqk, int R, int N, int P, int Na, void* stream

void mamba3_mimo_step_launch_f32(STEP_ARGS);
void mamba3_mimo_step_launch_bf16(STEP_ARGS);

static const void* vptr(const c10::optional<at::Tensor>& t, at::ScalarType st) {
    if (!t.has_value() || !t->defined()) return nullptr;
    TORCH_CHECK(t->is_cuda() && t->scalar_type() == st && t->is_contiguous(),
                "optional tensors must be contiguous CUDA tensors of the input dtype");
    return t->data_ptr();
}

void mamba3_cumulative_angles_launch(
    const float* angles, const float* dt, float* out, float* seg_tot,
    int B, int S, int H, int Na, int nseg, void* stream);
int mamba3_cumulative_angles_segments(int B, int S, int H);

torch::Tensor mamba3_cumulative_angles(torch::Tensor angles, torch::Tensor dt)
{
    M3_F32(angles); M3_F32(dt);
    TORCH_CHECK(angles.dim() == 4 && dt.dim() == 3, "angles is (B,S,H,Na), dt is (B,H,S)");
    const int B = (int)angles.size(0), S = (int)angles.size(1);
    const int H = (int)angles.size(2), Na = (int)angles.size(3);
    TORCH_CHECK(dt.size(0) == B && dt.size(1) == H && dt.size(2) == S,
                "dt must be (B,H,S) matching angles");
    auto out = torch::empty_like(angles);
    const int nseg = mamba3_cumulative_angles_segments(B, S, H);
    auto seg_tot = torch::empty({B * H, nseg, Na}, angles.options());
    mamba3_cumulative_angles_launch(
        angles.data_ptr<float>(), dt.data_ptr<float>(), out.data_ptr<float>(),
        seg_tot.data_ptr<float>(), B, S, H, Na, nseg,
        c10::cuda::getCurrentCUDAStream().stream());
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}

// Takes and returns no tensors, so it carries no meta kernel and belongs
// outside a compiled region; it exists so a caller, and the dispatch test, can
// ask which path a geometry takes without restating the gates.
std::vector<int64_t> mamba3_fwd_dispatch(
    int64_t mimo_rank, int64_t dstate, int64_t headdim, int64_t chunk_size, bool bf16)
{
    TORCH_CHECK(chunk_size > 0, "chunk_size must be positive");
    int state = 0, scan = 0, tf32 = 0;
    mamba3_fwd_dispatch_paths((int)mimo_rank, (int)dstate, (int)headdim,
                              (int)chunk_size, bf16 ? 1 : 0, &state, &scan, &tf32);
    return {state, scan, tf32};
}

torch::Tensor mamba3_mimo_step(
    torch::Tensor q, torch::Tensor k, torch::Tensor v, c10::optional<at::Tensor> z,
    torch::Tensor adt, torch::Tensor dt, torch::Tensor trap,
    torch::Tensor q_bias, torch::Tensor k_bias, torch::Tensor angles,
    torch::Tensor mimo_v, torch::Tensor mimo_o,
    c10::optional<at::Tensor> mimo_z, c10::optional<at::Tensor> D,
    torch::Tensor angle_in, torch::Tensor angle_out, torch::Tensor S_st,
    torch::Tensor kprev_in, torch::Tensor vprev_in,
    torch::Tensor kprev_out, torch::Tensor vprev_out)
{
    const auto sdt = q.scalar_type();
    TORCH_CHECK(sdt == at::kFloat || sdt == at::kBFloat16,
                "inputs must be float32 or bfloat16");
    for (const auto& t : {q, k, v, q_bias, k_bias, mimo_v, mimo_o})
        TORCH_CHECK(t.is_cuda() && t.is_contiguous() && t.scalar_type() == sdt,
                    "all value inputs must be contiguous CUDA tensors of one dtype");
    M3_F32(adt); M3_F32(dt); M3_F32(trap); M3_F32(angles);
    M3_F32(angle_in); M3_F32(angle_out); M3_F32(S_st);
    M3_F32(kprev_in); M3_F32(vprev_in); M3_F32(kprev_out); M3_F32(vprev_out);

    const int B = (int)q.size(0), R = (int)q.size(1), Gqk = (int)q.size(2), N = (int)q.size(3);
    const int H = (int)v.size(1), P = (int)v.size(2), Na = (int)angles.size(2);

    TORCH_CHECK(N % 2 == 0, "state dimension N must be even");
    TORCH_CHECK(H % Gqk == 0, "nheads must be divisible by the query-group count");
    TORCH_CHECK(Na <= N / 2, "angle count must not exceed N/2");
    TORCH_CHECK(R <= 8, "mimo_rank above 8 is not supported");
    TORCH_CHECK(S_st.size(2) == P && S_st.size(3) == N, "state shape does not match inputs");

    auto y = torch::empty({B, H, P}, v.options());
    auto fn = (sdt == at::kFloat) ? mamba3_mimo_step_launch_f32 : mamba3_mimo_step_launch_bf16;
    fn(q.data_ptr(), k.data_ptr(), v.data_ptr(), vptr(z, sdt),
       adt.data_ptr<float>(), dt.data_ptr<float>(), trap.data_ptr<float>(),
       q_bias.data_ptr(), k_bias.data_ptr(), angles.data_ptr<float>(),
       mimo_v.data_ptr(), mimo_o.data_ptr(), vptr(mimo_z, sdt), vptr(D, sdt),
       angle_in.data_ptr<float>(), angle_out.data_ptr<float>(), S_st.data_ptr<float>(),
       kprev_in.data_ptr<float>(), vprev_in.data_ptr<float>(),
       kprev_out.data_ptr<float>(), vprev_out.data_ptr<float>(), y.data_ptr(),
       B, H, Gqk, R, N, P, Na, c10::cuda::getCurrentCUDAStream().stream());
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}

#define FWD_ARGS \
    const void* q, const void* k, const void* v, const void* z, const void* q_bias, \
    const void* k_bias, const void* mimo_v, const void* mimo_o, const void* mimo_z, \
    const void* Dvec, const void* angles, const float* dt, const float* trap, \
    const float* dA_cs, const float* dA_cs_rev, const void* norm_w, float norm_eps, \
    int fused_norm, void* kv_ws, void* out, \
    int B, int S, int H, int Gqk, int R, int N, int P, int Na, int C, void* stream

void mamba3_mimo_fwd_launch_f32(FWD_ARGS);
void mamba3_mimo_fwd_launch_bf16(FWD_ARGS);

torch::Tensor mamba3_mimo_fwd(
    torch::Tensor q, torch::Tensor k, torch::Tensor v, c10::optional<at::Tensor> z,
    torch::Tensor q_bias, torch::Tensor k_bias, torch::Tensor mimo_v, torch::Tensor mimo_o,
    c10::optional<at::Tensor> mimo_z, c10::optional<at::Tensor> D,
    torch::Tensor angles, torch::Tensor dt, torch::Tensor trap,
    torch::Tensor dA_cs, torch::Tensor dA_cs_rev, int64_t chunk_size,
    c10::optional<at::Tensor> norm_weight, double norm_eps)
{
    const auto sdt = q.scalar_type();
    TORCH_CHECK(sdt == at::kFloat || sdt == at::kBFloat16,
                "inputs must be float32 or bfloat16");
    for (const auto& t : {q, k, v, q_bias, k_bias, mimo_v, mimo_o})
        TORCH_CHECK(t.is_cuda() && t.is_contiguous() && t.scalar_type() == sdt,
                    "all value inputs must be contiguous CUDA tensors of one dtype");
    // The schedule and the accumulated rotation angles stay float32 in both paths.
    M3_F32(dt); M3_F32(trap); M3_F32(dA_cs); M3_F32(dA_cs_rev); M3_F32(angles);

    const int B = (int)q.size(0), S = (int)q.size(1), R = (int)q.size(2);
    const int Gqk = (int)q.size(3), N = (int)q.size(4);
    const int H = (int)v.size(2), P = (int)v.size(3), Na = (int)angles.size(3);
    const int C = (int)chunk_size;
    TORCH_CHECK(C > 0, "chunk_size must be positive");
    const int Nc = (S + C - 1) / C;

    TORCH_CHECK(N % 2 == 0, "state dimension N must be even");
    TORCH_CHECK(H % Gqk == 0, "nheads must be divisible by the query-group count");
    TORCH_CHECK(Na <= N / 2, "angle count must not exceed N/2");
    // The tensor-core scan is instantiated for a fixed set of ranks; above it the
    // dispatch has no kernel to fall back to that carries the rank at runtime.
    TORCH_CHECK(R <= 8, "mimo_rank above 8 is not supported");

    const size_t smem = mamba3_fwd_smem_bytes(R, N, P, C, sdt == at::kBFloat16 ? 1 : 0);
    const size_t cap = (size_t)mamba3_max_smem_optin();
    TORCH_CHECK(smem <= cap,
                "chunk_size ", C, " with mimo_rank ", R, ", headdim ", P, " and state ", N,
                " needs ", smem / 1024, " KB of shared memory per block, above this "
                "device's ", cap / 1024, " KB opt-in limit");

    const int fused = (norm_weight.has_value() && norm_weight->defined()) ? 1 : 0;
    auto out = torch::empty({B, S, H, P}, v.options());
    // The state workspace carries the value dtype. chunk_scan casts the states
    // to bfloat16 the moment it reads them, so bf16 storage loses nothing there
    // and halves the traffic of the scan that walks this buffer.
    auto kv_ws = torch::empty({B, H, Nc, N, P}, v.options());

    auto fn = (sdt == at::kFloat) ? mamba3_mimo_fwd_launch_f32 : mamba3_mimo_fwd_launch_bf16;
    fn(q.data_ptr(), k.data_ptr(), v.data_ptr(), vptr(z, sdt), q_bias.data_ptr(),
       k_bias.data_ptr(), mimo_v.data_ptr(), mimo_o.data_ptr(), vptr(mimo_z, sdt), vptr(D, sdt),
       angles.data_ptr(), dt.data_ptr<float>(), trap.data_ptr<float>(),
       dA_cs.data_ptr<float>(), dA_cs_rev.data_ptr<float>(), vptr(norm_weight, sdt),
       (float)norm_eps, fused, kv_ws.data_ptr(), out.data_ptr(),
       B, S, H, Gqk, R, N, P, Na, C, c10::cuda::getCurrentCUDAStream().stream());
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}

#define BWD_ARGS \
    const void* q, const void* k, const void* v, const void* z, const void* q_bias, \
    const void* k_bias, const void* mimo_v, const void* mimo_o, const void* mimo_z, \
    const void* Dvec, const float* angles, const float* dt, const float* trap, \
    const float* dA_cs, const float* dA_cs_rev, const float* states, const void* norm_w, \
    float norm_eps, int fused_norm, const void* dy, float* dO_ws, float* dSt_ws, float* ws, \
    float* dq, float* dk, float* dv, float* dz, float* dq_bias, float* dk_bias, \
    float* dmimo_v, float* dmimo_o, float* dmimo_z, float* dD, float* dnorm_w, \
    float* d_dA_cs, float* d_dA_cs_rev, float* ddt, float* dtrap, float* dang, \
    int B, int S, int H, int Gqk, int R, int N, int P, int Na, int C, void* stream

void mamba3_mimo_bwd_launch_f32(BWD_ARGS);
void mamba3_mimo_bwd_launch_bf16(BWD_ARGS);
void mamba3_mimo_states_launch_f32(
    const void* k, const void* v, const void* k_bias, const void* mimo_v,
    const float* angles, const float* dt, const float* trap,
    const float* dA_cs, const float* dA_cs_rev, float* kv_ws,
    int B, int S, int H, int Gqk, int R, int N, int P, int Na, int C, void* stream);

std::vector<at::Tensor> mamba3_mimo_bwd(
    torch::Tensor q, torch::Tensor k, torch::Tensor v, c10::optional<at::Tensor> z,
    torch::Tensor q_bias, torch::Tensor k_bias, torch::Tensor mimo_v, torch::Tensor mimo_o,
    c10::optional<at::Tensor> mimo_z, c10::optional<at::Tensor> D,
    torch::Tensor angles, torch::Tensor dt, torch::Tensor trap,
    torch::Tensor dA_cs, torch::Tensor dA_cs_rev, int64_t chunk_size,
    c10::optional<at::Tensor> norm_weight, double norm_eps, torch::Tensor dy)
{
    const auto sdt = q.scalar_type();
    TORCH_CHECK(sdt == at::kFloat || sdt == at::kBFloat16,
                "backward inputs must be float32 or bfloat16");
    // The kernels index these with contiguous strides; a strided view would be
    // read as if it were packed and would return wrong gradients silently.
    for (const auto& t : {q, k, v, q_bias, k_bias, mimo_v, mimo_o})
        TORCH_CHECK(t.is_cuda() && t.is_contiguous() && t.scalar_type() == sdt,
                    "all value inputs must be contiguous CUDA tensors of one dtype");
    M3_F32(dt); M3_F32(trap); M3_F32(dA_cs); M3_F32(dA_cs_rev); M3_F32(angles);

    const int B = (int)q.size(0), S = (int)q.size(1), R = (int)q.size(2);
    const int Gqk = (int)q.size(3), N = (int)q.size(4);
    const int H = (int)v.size(2), P = (int)v.size(3), Na = (int)angles.size(3);
    TORCH_CHECK(chunk_size > 0, "chunk_size must be positive");
    TORCH_CHECK(R <= 8, "mimo_rank above 8 is not supported");
    const int C = (int)chunk_size, Nc = (S + C - 1) / C, CR = C * R;

    // The backward has no row-block streamed path, so its tiles are resident and
    // its budget is tighter than the forward's at the same geometry.
    const size_t smem = std::max(mamba3_bwd_smem_bytes(R, N, P, C),
                                 mamba3_states_smem_bytes(R, N, P, C));
    const size_t cap = (size_t)mamba3_max_smem_optin();
    TORCH_CHECK(smem <= cap,
                "backward with chunk_size ", C, ", mimo_rank ", R, ", headdim ", P,
                " and state ", N, " needs ", smem / 1024, " KB of shared memory per "
                "block, above this device's ", cap / 1024, " KB opt-in limit; reduce "
                "chunk_size * mimo_rank or headdim");

    auto f32 = v.options().dtype(at::kFloat);
    // The state recompute runs in float32 regardless; it feeds the inter term.
    auto kf = k.to(at::kFloat).contiguous(), vf = v.to(at::kFloat).contiguous();
    auto kbf = k_bias.to(at::kFloat).contiguous(), mvf = mimo_v.to(at::kFloat).contiguous();
    auto states = torch::empty({B, H, Nc, N, P}, f32);
    mamba3_mimo_states_launch_f32(
        kf.data_ptr(), vf.data_ptr(), kbf.data_ptr(), mvf.data_ptr(),
        angles.data_ptr<float>(), dt.data_ptr<float>(), trap.data_ptr<float>(),
        dA_cs.data_ptr<float>(), dA_cs_rev.data_ptr<float>(), states.data_ptr<float>(),
        B, S, H, Gqk, R, N, P, Na, C, c10::cuda::getCurrentCUDAStream().stream());
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    auto dO_ws  = torch::empty({B, H, Nc, CR, P}, f32);
    auto dSt_ws = torch::empty({B, H, Nc, N, P}, f32);
    // Per-block slots for the parameter gradients, reduced in block order so the
    // result does not depend on scheduling.
    const int64_t nblk = (int64_t)B * Nc;
    const int64_t gy = ((int64_t)N * P + 255) / 256;
    const int64_t nws = nblk * ((int64_t)3 * H * R * P + 2 * H * R * N + H * P + H)
                        + 2LL * B * H * S + (int64_t)B * H * Nc * gy;
    auto ws = torch::empty({nws}, f32);

    auto dq = torch::zeros_like(q, f32), dk = torch::zeros_like(k, f32);
    auto dv = torch::zeros_like(v, f32), dz = torch::zeros_like(v, f32);
    auto dqb = torch::zeros_like(q_bias, f32), dkb = torch::zeros_like(k_bias, f32);
    auto dmv = torch::zeros_like(mimo_v, f32), dmo = torch::zeros_like(mimo_o, f32);
    auto dmz = torch::zeros_like(mimo_v, f32);
    auto dD = torch::zeros({H}, f32), dnw = torch::zeros({H, P}, f32);
    auto dacs = torch::zeros_like(dA_cs), dacsr = torch::zeros_like(dA_cs_rev);
    auto ddt = torch::zeros_like(dt), dtrap = torch::zeros_like(trap);
    auto dang = torch::zeros_like(angles);

    const int fused = (norm_weight.has_value() && norm_weight->defined()) ? 1 : 0;
    auto bwd = (sdt == at::kFloat) ? mamba3_mimo_bwd_launch_f32 : mamba3_mimo_bwd_launch_bf16;
    bwd(q.data_ptr(), k.data_ptr(), v.data_ptr(), vptr(z, sdt), q_bias.data_ptr(),
        k_bias.data_ptr(), mimo_v.data_ptr(), mimo_o.data_ptr(), vptr(mimo_z, sdt),
        vptr(D, sdt), angles.data_ptr<float>(), dt.data_ptr<float>(),
        trap.data_ptr<float>(), dA_cs.data_ptr<float>(), dA_cs_rev.data_ptr<float>(),
        states.data_ptr<float>(), vptr(norm_weight, sdt), (float)norm_eps, fused,
        dy.contiguous().data_ptr(), dO_ws.data_ptr<float>(), dSt_ws.data_ptr<float>(),
        ws.data_ptr<float>(),
        dq.data_ptr<float>(), dk.data_ptr<float>(), dv.data_ptr<float>(),
        dz.data_ptr<float>(), dqb.data_ptr<float>(), dkb.data_ptr<float>(),
        dmv.data_ptr<float>(), dmo.data_ptr<float>(), dmz.data_ptr<float>(),
        dD.data_ptr<float>(), dnw.data_ptr<float>(), dacs.data_ptr<float>(),
        dacsr.data_ptr<float>(), ddt.data_ptr<float>(), dtrap.data_ptr<float>(),
        dang.data_ptr<float>(),
        B, S, H, Gqk, R, N, P, Na, C, c10::cuda::getCurrentCUDAStream().stream());
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return {dq, dk, dv, dz, dqb, dkb, dmv, dmo, dmz, dD, dnw, dacs, dacsr, ddt, dtrap, dang};
}

// registration.h comes from the kernel builder; a local JIT build has neither
// it nor TORCH_LIBRARY_EXPAND, and reaches the ops through torch.ops instead.
#if defined(__has_include) && __has_include("registration.h")
#  include "registration.h"
#  define MAMBA3_HAVE_REGISTRATION 1
#else
#  define TORCH_LIBRARY_EXPAND(NAME, MOD) TORCH_LIBRARY(NAME, MOD)
#  define CUDA_KERNEL 1
#endif

TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
    // Defined with its implementation: no tensor arguments, so there is no
    // dispatch key to select on and no meta kernel to register.
    ops.def("mamba3_fwd_dispatch(int mimo_rank, int dstate, int headdim, "
            "int chunk_size, bool bf16) -> int[]", &mamba3_fwd_dispatch);
    ops.def("mamba3_cumulative_angles(Tensor angles, Tensor dt) -> Tensor");
    ops.def("mamba3_mimo_step(Tensor q, Tensor k, Tensor v, Tensor? z, Tensor adt, "
            "Tensor dt, Tensor trap, Tensor q_bias, Tensor k_bias, Tensor angles, "
            "Tensor mimo_v, Tensor mimo_o, Tensor? mimo_z, Tensor? D, "
            "Tensor angle_in, Tensor(a!) angle_out, Tensor(b!) S_st, "
            "Tensor kprev_in, Tensor vprev_in, Tensor(c!) kprev_out, "
            "Tensor(d!) vprev_out) -> Tensor");
    ops.def("mamba3_mimo_fwd(Tensor q, Tensor k, Tensor v, Tensor? z, Tensor q_bias, "
            "Tensor k_bias, Tensor mimo_v, Tensor mimo_o, Tensor? mimo_z, Tensor? D, "
            "Tensor angles, Tensor dt, Tensor trap, Tensor dA_cs, Tensor dA_cs_rev, "
            "int chunk_size, Tensor? norm_weight, float norm_eps) -> Tensor");

    ops.def("mamba3_mimo_bwd(Tensor q, Tensor k, Tensor v, Tensor? z, Tensor q_bias, "
            "Tensor k_bias, Tensor mimo_v, Tensor mimo_o, Tensor? mimo_z, Tensor? D, "
            "Tensor angles, Tensor dt, Tensor trap, Tensor dA_cs, Tensor dA_cs_rev, "
            "int chunk_size, Tensor? norm_weight, float norm_eps, Tensor dy) -> Tensor[]");

#if defined(CUDA_KERNEL) || defined(ROCM_KERNEL)
    ops.impl("mamba3_cumulative_angles", torch::kCUDA, &mamba3_cumulative_angles);
    ops.impl("mamba3_mimo_step", torch::kCUDA, &mamba3_mimo_step);
    ops.impl("mamba3_mimo_fwd", torch::kCUDA, &mamba3_mimo_fwd);
    ops.impl("mamba3_mimo_bwd", torch::kCUDA, &mamba3_mimo_bwd);
#endif
}

#ifdef MAMBA3_HAVE_REGISTRATION
REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
#endif