"""DeepSeek NSA-inspired sparse attention — fused CUDA kernel (SM120).

Bench semantics match reference.nsa_attend:
  block importance = mean of (q·k / sqrt(D)) over causal keys in a block
                   = q · mean(k_block) / sqrt(D)
  top_n_blocks=8 union sliding_window=64; softmax over the union.
"""
from __future__ import annotations

from pathlib import Path

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

BLOCK_SIZE = 64
TOP_N_BLOCKS = 8
SLIDING_WINDOW = 64

_CPP_SRC = r"""
void nsa_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o);
"""

_mod = None


def _cuda_source() -> str:
    cu = Path(__file__).resolve().parent / "nsa_kernels.cu"
    return cu.read_text()


def _ext():
    global _mod
    if _mod is None:
        build_dir = Path("/tmp/nsa_cuda_build_v8")
        build_dir.mkdir(parents=True, exist_ok=True)
        _mod = load_inline(
            name="nsa_cuda_fwd_v8",
            cpp_sources=[_CPP_SRC],
            cuda_sources=[_cuda_source()],
            functions=["nsa_forward"],
            extra_cuda_cflags=[
                "-O3",
                "--use_fast_math",
                "-std=c++17",
                "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
                "-U__CUDA_NO_BFLOAT16_OPERATORS__",
            ],
            extra_ldflags=["-lcuda"],
            verbose=False,
            with_cuda=True,
            build_directory=str(build_dir),
        )
    return _mod


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:
        o = torch.empty_like(q)
        _ext().nsa_forward(q, k, v, o)
        return o


# ==================================================================
# ===== sidecar: nsa_kernels.cu (10493 bytes, loaded by solution.py) =====
# ==================================================================

#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_bf16.h>
#include <cuda_runtime.h>
#include <algorithm>
#include <cmath>

constexpr int kBlock = 64;
constexpr int kTopN = 8;
constexpr int kWindow = 64;

__device__ __forceinline__ float warp_sum(float v) {
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1)
        v += __shfl_xor_sync(0xffffffffu, v, off);
    return v;
}

__device__ __forceinline__ void load_bf16_pair(const __nv_bfloat16* p, float& a, float& b) {
    const __nv_bfloat162 v = __ldg(reinterpret_cast<const __nv_bfloat162*>(p));
    a = __bfloat162float(v.x);
    b = __bfloat162float(v.y);
}

__device__ __forceinline__ void insert_topk(float* ss, int* ii, float s, int idx) {
    const bool better =
        (ii[kTopN - 1] < 0) ||
        (s > ss[kTopN - 1]) ||
        (s == ss[kTopN - 1] && idx > ii[kTopN - 1]);
    if (!better) return;
    ss[kTopN - 1] = s;
    ii[kTopN - 1] = idx;
    #pragma unroll
    for (int k = kTopN - 1; k > 0; --k) {
        const bool swap =
            (ii[k - 1] < 0) ||
            (ss[k] > ss[k - 1]) ||
            (ss[k] == ss[k - 1] && ii[k] > ii[k - 1]);
        if (!swap) break;
        const float ts = ss[k - 1];
        ss[k - 1] = ss[k];
        ss[k] = ts;
        const int ti = ii[k - 1];
        ii[k - 1] = ii[k];
        ii[k] = ti;
    }
}

template <int D>
__device__ __forceinline__ float dot_key(
    const __nv_bfloat16* krow, const float* qv, int lane
) {
    constexpr int NP = D / 64;
    float dot = 0.f;
    #pragma unroll
    for (int i = 0; i < NP; ++i) {
        float a, b;
        load_bf16_pair(krow + i * 64 + 2 * lane, a, b);
        dot += qv[2 * i] * a + qv[2 * i + 1] * b;
    }
    return warp_sum(dot);
}

template <int D>
__device__ __forceinline__ void acc_value(
    const __nv_bfloat16* vrow, float p, float* acc, int lane
) {
    constexpr int NP = D / 64;
    #pragma unroll
    for (int i = 0; i < NP; ++i) {
        float a, b;
        load_bf16_pair(vrow + i * 64 + 2 * lane, a, b);
        acc[2 * i] += p * a;
        acc[2 * i + 1] += p * b;
    }
}

template <int D>
__device__ __forceinline__ void attend_rows(
    const __nv_bfloat16* kbase,
    const __nv_bfloat16* vbase,
    int n,
    const float* qv,
    float scale,
    float& m,
    float& lse,
    float* acc,
    int lane
) {
    constexpr int NE = D / 32;
    int j = 0;
    while (j < n) {
        const int take = min(4, n - j);
        float sc[4];
        #pragma unroll
        for (int u = 0; u < 4; ++u) {
            if (u < take)
                sc[u] = dot_key<D>(kbase + (j + u) * D, qv, lane) * scale;
            else
                sc[u] = -INFINITY;
        }
        const float bm = fmaxf(fmaxf(sc[0], sc[1]), fmaxf(sc[2], sc[3]));
        const float m_new = fmaxf(m, bm);
        const float alpha = __expf(m - m_new);
        lse *= alpha;
        #pragma unroll
        for (int i = 0; i < NE; ++i) acc[i] *= alpha;
        #pragma unroll
        for (int u = 0; u < 4; ++u) {
            if (u < take) {
                const float p = __expf(sc[u] - m_new);
                lse += p;
                acc_value<D>(vbase + (j + u) * D, p, acc, lane);
            }
        }
        m = m_new;
        j += take;
    }
}

template <int D>
__global__ void compress_k_kernel(
    const __nv_bfloat16* __restrict__ K,
    float* __restrict__ Ksum,
    int S,
    int n_blocks
) {
    const int bi = blockIdx.x;
    const int bh = blockIdx.y;
    const int s0 = bi * kBlock;
    const int len = min(kBlock, S - s0);
    const __nv_bfloat16* kb = K + (static_cast<size_t>(bh) * S + s0) * D;
    float* ob = Ksum + (static_cast<size_t>(bh) * n_blocks + bi) * D;
    for (int d = threadIdx.x; d < D; d += blockDim.x) {
        float sum = 0.f;
        for (int i = 0; i < len; ++i)
            sum += __bfloat162float(__ldg(kb + i * D + d));
        ob[d] = sum;
    }
}

// One warp = one query. Score blocks via compressed K, top-8, union
// with the causal sliding window, online-softmax over the unique keys.
template <int D, int WARPS>
__global__ void __launch_bounds__(WARPS * 32, 8)
nsa_query_kernel(
    const __nv_bfloat16* __restrict__ Q,
    const __nv_bfloat16* __restrict__ K,
    const __nv_bfloat16* __restrict__ V,
    const float* __restrict__ Ksum,
    __nv_bfloat16* __restrict__ O,
    int S,
    int n_blocks,
    int total_queries
) {
    constexpr int NP = D / 64;
    constexpr int NE = D / 32;
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    const int query_idx = blockIdx.x * WARPS + warp;
    if (query_idx >= total_queries) return;

    const int bh = query_idx / S;
    const int t = query_idx - bh * S;
    const float scale = rsqrtf(static_cast<float>(D));

    const __nv_bfloat16* qptr = Q + (static_cast<size_t>(bh) * S + t) * D;
    const __nv_bfloat16* khead = K + static_cast<size_t>(bh) * S * D;
    const __nv_bfloat16* vhead = V + static_cast<size_t>(bh) * S * D;
    const float* cbase = Ksum + static_cast<size_t>(bh) * n_blocks * D;

    float qv[NE];
    #pragma unroll
    for (int i = 0; i < NP; ++i)
        load_bf16_pair(qptr + i * 64 + 2 * lane, qv[2 * i], qv[2 * i + 1]);

    float top_s[kTopN];
    int top_i[kTopN];
    #pragma unroll
    for (int i = 0; i < kTopN; ++i) {
        top_s[i] = -1e30f;
        top_i[i] = -1;
    }

    const int curr = t >> 6;
    const float inv64 = scale * (1.f / 64.f);
    for (int bi = 0; bi < curr; ++bi) {
        float dot = 0.f;
        #pragma unroll
        for (int i = 0; i < NP; ++i) {
            const float* cp = cbase + bi * D + i * 64 + 2 * lane;
            dot += qv[2 * i] * __ldg(cp) + qv[2 * i + 1] * __ldg(cp + 1);
        }
        insert_topk(top_s, top_i, warp_sum(dot) * inv64, bi);
    }
    {
        const int s0 = curr * kBlock;
        const int len = t - s0 + 1;
        const float invn = scale / static_cast<float>(len);
        float dot = 0.f;
        #pragma unroll
        for (int i = 0; i < NP; ++i) {
            float sa = 0.f, sb = 0.f;
            const int od = i * 64 + 2 * lane;
            for (int j = 0; j < len; ++j) {
                float a, b;
                load_bf16_pair(khead + (static_cast<size_t>(s0 + j) * D) + od, a, b);
                sa += a;
                sb += b;
            }
            dot += qv[2 * i] * sa + qv[2 * i + 1] * sb;
        }
        insert_topk(top_s, top_i, warp_sum(dot) * invn, curr);
    }

    int sel[kTopN];
    int nsel = 0;
    bool take_curr = false, take_prev = false;
    const int prev = curr - 1;
    #pragma unroll
    for (int i = 0; i < kTopN; ++i) {
        const int bi = top_i[i];
        if (bi < 0) continue;
        sel[nsel++] = bi;
        if (bi == curr) take_curr = true;
        if (bi == prev) take_prev = true;
    }
    // Stream selected blocks in index order for sequential KV access.
    #pragma unroll
    for (int a = 0; a < kTopN - 1; ++a)
        #pragma unroll
        for (int b = a + 1; b < kTopN; ++b)
            if (b < nsel && a < nsel && sel[b] < sel[a]) {
                const int tmp = sel[a];
                sel[a] = sel[b];
                sel[b] = tmp;
            }

    float m = -INFINITY, lse = 0.f, acc[NE];
    #pragma unroll
    for (int i = 0; i < NE; ++i) acc[i] = 0.f;

    for (int i = 0; i < nsel; ++i) {
        const int bi = sel[i];
        const int s0 = bi * kBlock;
        const int n = min(kBlock, t + 1 - s0);
        attend_rows<D>(khead + s0 * D, vhead + s0 * D, n, qv, scale, m, lse, acc, lane);
    }
    if (!take_curr)
        attend_rows<D>(khead + curr * kBlock * D, vhead + curr * kBlock * D,
                      t - curr * kBlock + 1, qv, scale, m, lse, acc, lane);
    if (prev >= 0 && !take_prev) {
        const int w0 = max(0, t + 1 - kWindow);
        attend_rows<D>(khead + w0 * D, vhead + w0 * D, curr * kBlock - w0,
                      qv, scale, m, lse, acc, lane);
    }

    const float inv = 1.f / lse;
    __nv_bfloat16* optr = O + (static_cast<size_t>(bh) * S + t) * D;
    #pragma unroll
    for (int i = 0; i < NP; ++i) {
        const __nv_bfloat162 packed =
            __floats2bfloat162_rn(acc[2 * i] * inv, acc[2 * i + 1] * inv);
        *reinterpret_cast<__nv_bfloat162*>(optr + i * 64 + 2 * lane) = packed;
    }
}

void nsa_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o) {
    TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda() && o.is_cuda(), "cuda");
    TORCH_CHECK(q.scalar_type() == at::kBFloat16, "bf16");
    TORCH_CHECK(q.dim() == 4 && q.sizes() == k.sizes() && q.sizes() == v.sizes(), "shape");

    const auto q_c = q.contiguous();
    const auto k_c = k.contiguous();
    const auto v_c = v.contiguous();
    TORCH_CHECK(o.is_contiguous(), "o contig");

    const int B = static_cast<int>(q_c.size(0));
    const int H = static_cast<int>(q_c.size(1));
    const int S = static_cast<int>(q_c.size(2));
    const int D = static_cast<int>(q_c.size(3));
    TORCH_CHECK(D == 64 || D == 128, "D");

    const int n_blocks = (S + kBlock - 1) / kBlock;
    const int BH = B * H;
    const int total = BH * S;

    static thread_local at::Tensor ksum_cache;
    if (!ksum_cache.defined() || ksum_cache.size(0) != BH ||
        ksum_cache.size(1) != n_blocks || ksum_cache.size(2) != D ||
        ksum_cache.device() != q_c.device()) {
        ksum_cache = at::empty({BH, n_blocks, D}, q_c.options().dtype(at::kFloat));
    }

    const c10::cuda::CUDAGuard guard(q_c.device());
    auto stream = at::cuda::getCurrentCUDAStream();

    const auto* qp = reinterpret_cast<const __nv_bfloat16*>(q_c.data_ptr());
    const auto* kp = reinterpret_cast<const __nv_bfloat16*>(k_c.data_ptr());
    const auto* vp = reinterpret_cast<const __nv_bfloat16*>(v_c.data_ptr());
    auto* op = reinterpret_cast<__nv_bfloat16*>(o.data_ptr());
    float* cp = ksum_cache.data_ptr<float>();

    const dim3 cgrid(n_blocks, BH);
    constexpr int WARPS = 8;
    const int qgrid = (total + WARPS - 1) / WARPS;

    if (D == 64) {
        compress_k_kernel<64><<<cgrid, 128, 0, stream>>>(kp, cp, S, n_blocks);
        nsa_query_kernel<64, WARPS><<<qgrid, WARPS * 32, 0, stream>>>(
            qp, kp, vp, cp, op, S, n_blocks, total);
    } else {
        compress_k_kernel<128><<<cgrid, 128, 0, stream>>>(kp, cp, S, n_blocks);
        nsa_query_kernel<128, WARPS><<<qgrid, WARPS * 32, 0, stream>>>(
            qp, kp, vp, cp, op, S, n_blocks, total);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
}