import math

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

CUDA_SRC = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <math.h>

namespace {

constexpr int BLK = 64;
constexpr int TOPN = 8;
constexpr int NCHUNK = 2 * (TOPN + 1);
constexpr int WARPS = 8;
constexpr int MAXI = 8;
constexpr unsigned FULL = 0xffffffffu;

__device__ __forceinline__ float bf(const __nv_bfloat16 x) { return __bfloat162float(x); }

__device__ __forceinline__ float warp_sum(float v) {
    for (int o = 16; o > 0; o >>= 1) v += __shfl_xor_sync(FULL, v, o);
    return v;
}

__device__ __forceinline__ float warp_max(float v) {
    for (int o = 16; o > 0; o >>= 1) v = fmaxf(v, __shfl_xor_sync(FULL, v, o));
    return v;
}

template <int DPL>
__device__ __forceinline__ void load_lane(const __nv_bfloat16* row, int lane, float (&f)[DPL]) {
    if constexpr (DPL == 2) {
        const __nv_bfloat162 x = *reinterpret_cast<const __nv_bfloat162*>(row + 2 * lane);
        f[0] = __low2float(x);
        f[1] = __high2float(x);
    } else {
        const uint2 x = *reinterpret_cast<const uint2*>(row + 4 * lane);
        f[0] = __uint_as_float(x.x << 16);
        f[1] = __uint_as_float(x.x & 0xffff0000u);
        f[2] = __uint_as_float(x.y << 16);
        f[3] = __uint_as_float(x.y & 0xffff0000u);
    }
}

template <int DPL>
__device__ __forceinline__ float row_dot(const __nv_bfloat16* row, int lane, const float (&qv)[DPL]) {
    float kf[DPL];
    load_lane<DPL>(row, lane, kf);
    float s = 0.f;
    #pragma unroll
    for (int e = 0; e < DPL; ++e) s += qv[e] * kf[e];
    return s;
}

// v[i] is this lane's partial for row i; returns the complete sum for row `lane`.
__device__ __forceinline__ float reduce_rows32(float (&v)[32], int lane) {
    #pragma unroll
    for (int i = 0; i < 16; ++i) {
        const bool hi = lane & 16;
        const float send = hi ? v[i] : v[i + 16];
        const float keep = hi ? v[i + 16] : v[i];
        v[i] = keep + __shfl_xor_sync(FULL, send, 16);
    }
    #pragma unroll
    for (int i = 0; i < 8; ++i) {
        const bool hi = lane & 8;
        const float send = hi ? v[i] : v[i + 8];
        const float keep = hi ? v[i + 8] : v[i];
        v[i] = keep + __shfl_xor_sync(FULL, send, 8);
    }
    #pragma unroll
    for (int i = 0; i < 4; ++i) {
        const bool hi = lane & 4;
        const float send = hi ? v[i] : v[i + 4];
        const float keep = hi ? v[i + 4] : v[i];
        v[i] = keep + __shfl_xor_sync(FULL, send, 4);
    }
    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        const bool hi = lane & 2;
        const float send = hi ? v[i] : v[i + 2];
        const float keep = hi ? v[i + 2] : v[i];
        v[i] = keep + __shfl_xor_sync(FULL, send, 2);
    }
    const bool hi = lane & 1;
    const float send = hi ? v[0] : v[1];
    const float keep = hi ? v[1] : v[0];
    return keep + __shfl_xor_sync(FULL, send, 1);
}

__global__ void block_sum_kernel(const __nv_bfloat16* __restrict__ k, float* __restrict__ msum,
                                 int S, int nb, int D) {
    const int bi = blockIdx.x;
    const int bh = blockIdx.y;
    const int d = threadIdx.x;
    const __nv_bfloat16* kb = k + (size_t)bh * S * D;
    const int j0 = bi * BLK;
    const int j1 = min(j0 + BLK, S);
    float acc = 0.f;
    for (int j = j0; j < j1; ++j) acc += bf(kb[(size_t)j * D + d]);
    msum[((size_t)bh * nb + bi) * D + d] = acc;
}

template <int D>
__global__ void __launch_bounds__(WARPS * 32)
nsa_fwd_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ k,
               const __nv_bfloat16* __restrict__ v, const float* __restrict__ msum,
               __nv_bfloat16* __restrict__ out, int S, int nb, float scale) {
    constexpr int DP = D + 1;
    constexpr int DPL = D / 32;
    extern __shared__ float smem[];
    float* ms_s = smem;
    float* q_s = ms_s + (size_t)nb * DP;

    const int bh = blockIdx.y;
    const int warp = threadIdx.x >> 5;
    const int lane = threadIdx.x & 31;
    const int t_base = blockIdx.x * WARPS;
    const int t_max = min(t_base + WARPS - 1, S - 1);
    const int nb_need = min(nb, (t_max >> 6) + 1);

    for (int idx = threadIdx.x; idx < nb_need * D; idx += blockDim.x) {
        const int row = idx / D;
        const int d = idx - row * D;
        ms_s[row * DP + d] = msum[((size_t)bh * nb + row) * D + d];
    }
    __syncthreads();

    const int t = t_base + warp;
    const bool active = t < S;
    const int tt = active ? t : 0;
    const int tb = tt >> 6;
    const int dd = tt & 63;

    float* q_w = q_s + warp * D;
    const __nv_bfloat16* qr = q + ((size_t)bh * S + tt) * D;
    for (int d = lane; d < D; d += 32) q_w[d] = active ? bf(qr[d]) : 0.f;
    __syncwarp();

    float qv[DPL];
    #pragma unroll
    for (int e = 0; e < DPL; ++e) qv[e] = q_w[lane * DPL + e];

    const __nv_bfloat16* kb = k + (size_t)bh * S * D;
    const __nv_bfloat16* vb = v + (size_t)bh * S * D;

    float diag_part = 0.f;
    #pragma unroll
    for (int h = 0; h < 2; ++h) {
        float part[32];
        #pragma unroll
        for (int i = 0; i < 32; ++i) {
            const int j = min(tb * BLK + h * 32 + i, S - 1);
            part[i] = row_dot<DPL>(kb + (size_t)j * D, lane, qv);
        }
        const float s = reduce_rows32(part, lane);
        if (active && h * 32 + lane <= dd) diag_part += s * scale;
    }
    const float diag_imp = warp_sum(diag_part) / (float)(dd + 1);

    float imp[MAXI];
    #pragma unroll
    for (int i = 0; i < MAXI; ++i) {
        const int c = lane + 32 * i;
        if (c < tb) {
            const float* mr = ms_s + c * DP;
            float s = 0.f;
            for (int d = 0; d < D; ++d) s += q_w[d] * mr[d];
            imp[i] = s * (scale / (float)BLK);
        } else if (c == tb) {
            imp[i] = diag_imp;
        } else {
            imp[i] = -INFINITY;
        }
    }

    unsigned taken = 0u;
    int sel[TOPN];
    #pragma unroll
    for (int r = 0; r < TOPN; ++r) {
        float bv = -INFINITY;
        int bi = -1;
        #pragma unroll
        for (int i = 0; i < MAXI; ++i) {
            const int c = lane + 32 * i;
            const bool ok = (c <= tb) && !((taken >> i) & 1u);
            if (ok && (imp[i] > bv || (imp[i] == bv && c > bi))) {
                bv = imp[i];
                bi = c;
            }
        }
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) {
            const float ov = __shfl_xor_sync(FULL, bv, o);
            const int oi = __shfl_xor_sync(FULL, bi, o);
            if (ov > bv || (ov == bv && oi > bi)) {
                bv = ov;
                bi = oi;
            }
        }
        sel[r] = (bi >= 0 && bv > -INFINITY) ? bi : -1;
        if (sel[r] >= 0 && (bi & 31) == lane) taken |= 1u << (bi >> 5);
    }

    float sc[NCHUNK];
    #pragma unroll
    for (int c = 0; c < NCHUNK; ++c) {
        const int h = c & 1;
        int jbase;
        bool seg;
        if (c < 2 * TOPN) {
            const int b = sel[c >> 1];
            seg = (b >= 0) && active;
            jbase = b * BLK + h * 32;
        } else {
            seg = active;
            jbase = tt - (BLK - 1) + h * 32;
        }
        float part[32];
        #pragma unroll
        for (int i = 0; i < 32; ++i) {
            const int j = min(max(jbase + i, 0), S - 1);
            part[i] = row_dot<DPL>(kb + (size_t)j * D, lane, qv);
        }
        const float s = reduce_rows32(part, lane);
        const int key = jbase + lane;
        bool ok = seg && key >= 0 && key <= tt;
        if (c >= 2 * TOPN) {
            const int jb = key >> 6;
            #pragma unroll
            for (int r = 0; r < TOPN; ++r) ok = ok && (sel[r] != jb);
        }
        sc[c] = ok ? s * scale : -INFINITY;
    }

    float lmax = -INFINITY;
    #pragma unroll
    for (int c = 0; c < NCHUNK; ++c) lmax = fmaxf(lmax, sc[c]);
    const float m = warp_max(lmax);
    float pc[NCHUNK];
    float lsum = 0.f;
    #pragma unroll
    for (int c = 0; c < NCHUNK; ++c) {
        pc[c] = (sc[c] == -INFINITY) ? 0.f : expf(sc[c] - m);
        lsum += pc[c];
    }
    const float l = warp_sum(lsum);

    float acc[DPL];
    #pragma unroll
    for (int e = 0; e < DPL; ++e) acc[e] = 0.f;
    #pragma unroll
    for (int c = 0; c < NCHUNK; ++c) {
        const int h = c & 1;
        const int jbase = (c < 2 * TOPN) ? (sel[c >> 1] * BLK + h * 32) : (tt - (BLK - 1) + h * 32);
        for (int i = 0; i < 32; ++i) {
            const float p = __shfl_sync(FULL, pc[c], i);
            if (p != 0.f) {
                float vf[DPL];
                load_lane<DPL>(vb + (size_t)(jbase + i) * D, lane, vf);
                #pragma unroll
                for (int e = 0; e < DPL; ++e) acc[e] += p * vf[e];
            }
        }
    }

    if (active) {
        const float inv = 1.f / l;
        __nv_bfloat16* orow = out + ((size_t)bh * S + t) * D;
        #pragma unroll
        for (int e = 0; e < DPL; ++e) orow[lane * DPL + e] = __float2bfloat16(acc[e] * inv);
    }
}

}  // namespace

torch::Tensor nsa_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, double scale_d) {
    TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(q.scalar_type() == torch::kBFloat16 && k.scalar_type() == torch::kBFloat16 &&
                    v.scalar_type() == torch::kBFloat16,
                "inputs must be bfloat16");
    TORCH_CHECK(q.is_contiguous() && k.is_contiguous() && v.is_contiguous(), "inputs must be contiguous");
    TORCH_CHECK(q.dim() == 4, "inputs must be (B, H, S, D)");
    const int B = q.size(0), H = q.size(1), S = q.size(2), D = q.size(3);
    TORCH_CHECK(D == 64 || D == 128, "D must be 64 or 128");
    const int nb = (S + BLK - 1) / BLK;
    TORCH_CHECK(nb <= 256, "S too large for candidate slots");
    const int BH = B * H;

    auto msum = torch::empty({BH, nb, D}, q.options().dtype(torch::kFloat32));
    auto out = torch::empty_like(q);
    const float scale = (float)scale_d;
    auto stream = at::cuda::getCurrentCUDAStream();

    block_sum_kernel<<<dim3(nb, BH), D, 0, stream>>>(
        reinterpret_cast<const __nv_bfloat16*>(k.data_ptr<at::BFloat16>()),
        msum.data_ptr<float>(), S, nb, D);

    const size_t smem = ((size_t)nb * (D + 1) + WARPS * D) * sizeof(float);
    const dim3 grid((S + WARPS - 1) / WARPS, BH);
    const auto* qp = reinterpret_cast<const __nv_bfloat16*>(q.data_ptr<at::BFloat16>());
    const auto* kp = reinterpret_cast<const __nv_bfloat16*>(k.data_ptr<at::BFloat16>());
    const auto* vp = reinterpret_cast<const __nv_bfloat16*>(v.data_ptr<at::BFloat16>());
    auto* op = reinterpret_cast<__nv_bfloat16*>(out.data_ptr<at::BFloat16>());
    if (D == 64) {
        cudaFuncSetAttribute(nsa_fwd_kernel<64>, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem);
        nsa_fwd_kernel<64><<<grid, WARPS * 32, smem, stream>>>(qp, kp, vp, msum.data_ptr<float>(), op, S, nb, scale);
    } else {
        cudaFuncSetAttribute(nsa_fwd_kernel<128>, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem);
        nsa_fwd_kernel<128><<<grid, WARPS * 32, smem, stream>>>(qp, kp, vp, msum.data_ptr<float>(), op, S, nb, scale);
    }
    return out;
}
"""

CPP_SRC = "torch::Tensor nsa_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, double scale_d);"

_ext = load_inline(
    name="nsa_sparse_ext_v2",
    cpp_sources=CPP_SRC,
    cuda_sources=CUDA_SRC,
    functions=["nsa_forward"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)

B, H, S, D = 1, 16, 1024, 64


class Model(nn.Module):
    def __init__(self, B: int, H: int, S: int, D: int):
        super().__init__()
        self.B, self.H, self.S, self.D = B, H, S, D
        self.register_buffer("_dummy", torch.zeros(1, dtype=torch.bfloat16))

    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
        return _ext.nsa_forward(q.contiguous(), k.contiguous(), v.contiguous(), 1.0 / math.sqrt(self.D))


def get_init_inputs():
    return [B, H, S, D]


def get_inputs():
    q = torch.randn(B, H, S, D, dtype=torch.bfloat16)
    k = torch.randn(B, H, S, D, dtype=torch.bfloat16)
    v = torch.randn(B, H, S, D, dtype=torch.bfloat16)
    return [q, k, v]