"""CUDA implementation of the benchmark's block-selected causal attention."""
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline


_CPP = r"""
#include <torch/extension.h>
void nsa_cuda(torch::Tensor q, torch::Tensor k, torch::Tensor v,
              torch::Tensor prefix, torch::Tensor out);
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("run", &nsa_cuda); }
"""

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

__global__ void prefix_kernel(const __nv_bfloat16* __restrict__ k,
                              float* __restrict__ p, int S, int D, int NB) {
    int bi = blockIdx.x, bh = blockIdx.y;
    int d = threadIdx.x;
    if (d >= D) return;
    const __nv_bfloat16* kh = k + (size_t)bh * S * D;
    float* ph = p + (size_t)bh * S * D;
    int start = bi * 64;
    int end = min(start + 64, S);
    float sum = 0.0f;
    for (int j = start; j < end; ++j) {
        sum += __bfloat162float(kh[(size_t)j * D + d]);
        ph[(size_t)j * D + d] = sum / float(j - start + 1);
    }
}

__global__ void nsa_kernel(const __nv_bfloat16* __restrict__ q,
                           const __nv_bfloat16* __restrict__ k,
                           const __nv_bfloat16* __restrict__ v,
                           const float* __restrict__ prefix,
                           __nv_bfloat16* __restrict__ out,
                           int S, int D, int NB, float scale) {
    int t = blockIdx.x, bh = blockIdx.y;
    int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5;
    const __nv_bfloat16* qh = q + (size_t)bh * S * D;
    const __nv_bfloat16* kh = k + (size_t)bh * S * D;
    const __nv_bfloat16* vh = v + (size_t)bh * S * D;
    const float* ph = prefix + (size_t)bh * S * D;
    __nv_bfloat16* oh = out + (size_t)bh * S * D;
    const __nv_bfloat16* qt = qh + (size_t)t * D;
    __shared__ float importance[128];
    __shared__ int selected[8];
    __shared__ int indices[576];
    __shared__ float logits[576];
    __shared__ int count;
    __shared__ float max_logit, denom;

    int valid_blocks = t / 64 + 1;
    for (int bi = warp; bi < valid_blocks; bi += 4) {
        int last = bi == t / 64 ? t : min((bi + 1) * 64, S) - 1;
        const float* mean = ph + (size_t)last * D;
        float partial = 0.0f;
        for (int d = lane; d < D; d += 32)
            partial = fmaf(__bfloat162float(qt[d]), mean[d], partial);
        for (int delta = 16; delta; delta >>= 1)
            partial += __shfl_down_sync(0xffffffff, partial, delta);
        if (lane == 0) importance[bi] = partial * scale;
    }
    __syncthreads();

    if (tid == 0) {
        int nsel = min(valid_blocks, 8);
        for (int i = 0; i < nsel; ++i) {
            float best = -INFINITY;
            int best_b = -1;
            for (int bi = 0; bi < valid_blocks; ++bi) {
                float x = importance[bi];
                if (x > best || (x == best && bi > best_b)) {
                    best = x; best_b = bi;
                }
            }
            selected[i] = best_b;
            importance[best_b] = -INFINITY;
        }
        count = 0;
        for (int i = 0; i < nsel; ++i) {
            int start = selected[i] * 64;
            int end = min(start + 64, t + 1);
            for (int j = start; j < end; ++j) indices[count++] = j;
        }
        int w0 = max(0, t - 63);
        for (int j = w0; j <= t; ++j) {
            bool exists = false;
            for (int i = 0; i < nsel; ++i)
                if (selected[i] == j / 64) exists = true;
            if (!exists) indices[count++] = j;
        }
    }
    __syncthreads();

    for (int i = warp; i < count; i += 4) {
        const __nv_bfloat16* kj = kh + (size_t)indices[i] * D;
        float partial = 0.0f;
        for (int d = lane; d < D; d += 32)
            partial = fmaf(__bfloat162float(qt[d]), __bfloat162float(kj[d]), partial);
        for (int delta = 16; delta; delta >>= 1)
            partial += __shfl_down_sync(0xffffffff, partial, delta);
        if (lane == 0) logits[i] = partial * scale;
    }
    __syncthreads();
    if (tid == 0) {
        float m = -INFINITY;
        for (int i = 0; i < count; ++i) m = fmaxf(m, logits[i]);
        max_logit = m;
    }
    __syncthreads();
    for (int i = tid; i < count; i += blockDim.x)
        logits[i] = expf(logits[i] - max_logit);
    __syncthreads();
    if (tid == 0) {
        float z = 0.0f;
        for (int i = 0; i < count; ++i) z += logits[i];
        denom = z;
    }
    __syncthreads();
    for (int d = tid; d < D; d += blockDim.x) {
        float acc = 0.0f;
        for (int i = 0; i < count; ++i)
            acc = fmaf(logits[i], __bfloat162float(vh[(size_t)indices[i] * D + d]), acc);
        oh[(size_t)t * D + d] = __float2bfloat16(acc / denom);
    }
}

void nsa_cuda(torch::Tensor q, torch::Tensor k, torch::Tensor v,
              torch::Tensor prefix, torch::Tensor out) {
    int B = q.size(0), H = q.size(1), S = q.size(2), D = q.size(3);
    int NB = (S + 63) / 64;
    auto stream = at::cuda::getCurrentCUDAStream();
    prefix_kernel<<<dim3(NB, B * H), D, 0, stream>>>(
        (const __nv_bfloat16*)k.data_ptr(), (float*)prefix.data_ptr(), S, D, NB);
    nsa_kernel<<<dim3(S, B * H), 128, 0, stream>>>(
        (const __nv_bfloat16*)q.data_ptr(), (const __nv_bfloat16*)k.data_ptr(),
        (const __nv_bfloat16*)v.data_ptr(), (const float*)prefix.data_ptr(),
        (__nv_bfloat16*)out.data_ptr(), S, D, NB, rsqrtf(float(D)));
}
"""

_extension = None


def _get_extension():
    global _extension
    if _extension is None:
        _extension = load_inline(
            name="nsa_sparse_cuda_v1",
            cpp_sources=_CPP,
            cuda_sources=_CUDA,
            functions=None,
            extra_cflags=["-O3"],
            extra_cuda_cflags=["-O3", "-arch=sm_120"],
        )
    return _extension


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:
        out = torch.empty_like(q)
        prefix = torch.empty(q.shape, device=q.device, dtype=torch.float32)
        _get_extension().run(q, k, v, prefix, out)
        return out