"""DeepSeek NSA-inspired sparse attention — high-performance CUDA kernel.

Bench semantics matching reference.nsa_attend:
  block_size=64, top_n_blocks=8, sliding_window=64, causal, bf16 I/O.

Architecture:
  Query-tiled CTAs stream K through shared memory.
  Phase 1: parallel block-importance (4 lanes/query) + online top-8.
  Phase 2: warp-owned sparse online-softmax over selected ∪ window.
"""
from __future__ import annotations

from typing import Optional

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

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

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

static constexpr int BS   = 64;
static constexpr int TOPN = 8;
static constexpr int WIN  = 64;

__device__ __forceinline__ float bf2f(__nv_bfloat16 x) {
    return __bfloat162float(x);
}
__device__ __forceinline__ __nv_bfloat16 f2bf(float x) {
    return __float2bfloat16(x);
}

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

__device__ __forceinline__ float warp_sum_bcast(float v) {
    return __shfl_sync(0xffffffff, warp_sum(v), 0);
}

__device__ __forceinline__ void topk_insert(float* sc, int* ix, float s, int bi) {
    if (s < sc[TOPN - 1]) return;
    int pos = TOPN - 1;
    while (pos > 0 && s > sc[pos - 1]) {
        sc[pos] = sc[pos - 1];
        ix[pos] = ix[pos - 1];
        --pos;
    }
    sc[pos] = s;
    ix[pos] = bi;
}

template <int TD>
__device__ __forceinline__ float dot_ff(const float* __restrict__ q,
                                        const float* __restrict__ k) {
    float acc = 0.f;
    #pragma unroll
    for (int d = 0; d < TD; d += 4) {
        acc += q[d] * k[d] + q[d+1] * k[d+1] + q[d+2] * k[d+2] + q[d+3] * k[d+3];
    }
    return acc;
}

// =====================================================================
// TD head dim, QT query tile, NT threads per block
// =====================================================================
template <int TD, int QT, int NT>
__global__ void __launch_bounds__(NT, 4)
nsa_kernel(
    const __nv_bfloat16* __restrict__ Q,
    const __nv_bfloat16* __restrict__ K,
    const __nv_bfloat16* __restrict__ V,
    __nv_bfloat16* __restrict__ O,
    int B_, int H_, int S_, int D_
) {
    const int bh   = blockIdx.x;
    const int tile = blockIdx.y;
    const int t0   = tile * QT;
    if (t0 >= S_ || bh >= B_ * H_) return;

    const int tid  = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;
    constexpr int NW = NT / 32;

    const long long head_off = (static_cast<long long>(bh) * S_) * D_;
    const __nv_bfloat16* q_base = Q + head_off;
    const __nv_bfloat16* k_base = K + head_off;
    const __nv_bfloat16* v_base = V + head_off;
    __nv_bfloat16*       o_base = O + head_off;

    const float scale = rsqrtf(static_cast<float>(D_));
    const int n_blocks = (S_ + BS - 1) / BS;
    const int t_end = min(t0 + QT, S_);
    const int nq = t_end - t0;

    // smem: q_f[QT*TD] | k_f[BS*TD] | sel[QT*TOPN] | bsc[QT]
    extern __shared__ float smem[];
    float* q_f  = smem;
    float* k_f  = q_f + QT * TD;
    int*   sel  = reinterpret_cast<int*>(k_f + BS * TD);
    float* bsc  = reinterpret_cast<float*>(sel + QT * TOPN);

    // Load Q as float
    for (int i = tid; i < nq * TD; i += NT) {
        int qi = i / TD;
        int d  = i - qi * TD;
        q_f[qi * TD + d] = (d < D_)
            ? bf2f(q_base[(static_cast<long long>(t0 + qi)) * D_ + d])
            : 0.f;
    }
    for (int i = nq * TD + tid; i < QT * TD; i += NT) q_f[i] = 0.f;
    __syncthreads();

    float top_sc[TOPN];
    int   top_bi[TOPN];
    #pragma unroll
    for (int s = 0; s < TOPN; ++s) { top_sc[s] = -1e30f; top_bi[s] = -1; }

    // =================================================================
    // Phase 1
    // =================================================================
    constexpr int LPQ = 4;
    constexpr int QPW = 32 / LPQ; // 8 queries per warp

    for (int bi = 0; bi < n_blocks; ++bi) {
        const int s0 = bi * BS;
        if (s0 > t_end - 1) break;
        const int nk = min(BS, S_ - s0);

        // Load K block → float smem (coalesced)
        for (int i = tid; i < nk * D_; i += NT) {
            int j = i / D_;
            int d = i - j * D_;
            k_f[j * TD + d] = bf2f(k_base[(static_cast<long long>(s0 + j)) * D_ + d]);
        }
        for (int i = nk * TD + tid; i < BS * TD; i += NT) k_f[i] = 0.f;
        __syncthreads();

        // Score: 4 lanes/query, all lanes participate in shfl
        {
            const int qi_local = lane / LPQ;
            const int sub      = lane - qi_local * LPQ;
            const int qi       = warp * QPW + qi_local;
            const bool active  = (qi < nq);

            float local = 0.f;
            int j1 = 0;
            bool causal = false;
            if (active) {
                const int t = t0 + qi;
                if (s0 <= t) {
                    causal = true;
                    j1 = min(s0 + BS, t + 1) - s0;
                    for (int j = sub; j < j1; j += LPQ)
                        local += dot_ff<TD>(q_f + qi * TD, k_f + j * TD);
                }
            }
            local += __shfl_xor_sync(0xffffffff, local, 1);
            local += __shfl_xor_sync(0xffffffff, local, 2);
            if (active && sub == 0) {
                bsc[qi] = causal ? (local / static_cast<float>(j1)) * scale : -1e9f;
            }
        }
        __syncthreads();

        if (tid < nq) topk_insert(top_sc, top_bi, bsc[tid], bi);
        __syncthreads();
    }

    if (tid < nq) {
        #pragma unroll
        for (int s = 0; s < TOPN; ++s) sel[tid * TOPN + s] = top_bi[s];
    }
    __syncthreads();

    // =================================================================
    // Phase 2 — sparse online attention
    // =================================================================
    constexpr int EPT = (TD + 31) / 32;

    for (int qi = warp; qi < nq; qi += NW) {
        const int t = t0 + qi;
        const int w0 = (t + 1 > WIN) ? (t + 1 - WIN) : 0;

        int my_sel[TOPN];
        #pragma unroll
        for (int s = 0; s < TOPN; ++s) my_sel[s] = sel[qi * TOPN + s];

        // Selected-block bitset for n_blocks <= 256
        unsigned long long sel_bits0 = 0ull, sel_bits1 = 0ull, sel_bits2 = 0ull, sel_bits3 = 0ull;
        #pragma unroll
        for (int s = 0; s < TOPN; ++s) {
            int bi = my_sel[s];
            if (bi < 0) continue;
            if (bi < 64)       sel_bits0 |= (1ull << bi);
            else if (bi < 128) sel_bits1 |= (1ull << (bi - 64));
            else if (bi < 192) sel_bits2 |= (1ull << (bi - 128));
            else if (bi < 256) sel_bits3 |= (1ull << (bi - 192));
        }
        auto is_sel = [&](int bi) -> bool {
            if (bi < 64)       return (sel_bits0 >> bi) & 1ull;
            if (bi < 128)      return (sel_bits1 >> (bi - 64)) & 1ull;
            if (bi < 192)      return (sel_bits2 >> (bi - 128)) & 1ull;
            if (bi < 256)      return (sel_bits3 >> (bi - 192)) & 1ull;
            return false;
        };

        float q_reg[EPT];
        #pragma unroll
        for (int e = 0; e < EPT; ++e) {
            int d = lane + e * 32;
            q_reg[e] = (d < D_) ? q_f[qi * TD + d] : 0.f;
        }

        float m_i = -FLT_MAX;
        float l_i = 0.f;
        float acc[EPT];
        #pragma unroll
        for (int e = 0; e < EPT; ++e) acc[e] = 0.f;

        // Process one key from global memory (all lanes)
        #define PROCESS_KEY(j_expr)                                                        \
        do {                                                                               \
            const int j__ = (j_expr);                                                      \
            const __nv_bfloat16* kj = k_base + (static_cast<long long>(j__)) * D_;         \
            const __nv_bfloat16* vj = v_base + (static_cast<long long>(j__)) * D_;         \
            float partial = 0.f;                                                           \
            _Pragma("unroll")                                                              \
            for (int e = 0; e < EPT; ++e) {                                                \
                int d = lane + e * 32;                                                     \
                if (d < D_) partial += q_reg[e] * bf2f(kj[d]);                             \
            }                                                                              \
            float score = warp_sum_bcast(partial) * scale;                                 \
            float m_new = fmaxf(m_i, score);                                               \
            float ed = __expf(m_i - m_new);                                                \
            float es = __expf(score - m_new);                                              \
            _Pragma("unroll")                                                              \
            for (int e = 0; e < EPT; ++e) acc[e] *= ed;                                    \
            l_i = l_i * ed + es;                                                           \
            m_i = m_new;                                                                   \
            _Pragma("unroll")                                                              \
            for (int e = 0; e < EPT; ++e) {                                                \
                int d = lane + e * 32;                                                     \
                if (d < D_) acc[e] += es * bf2f(vj[d]);                                    \
            }                                                                              \
        } while (0)

        // Pass A: selected blocks
        #pragma unroll
        for (int s = 0; s < TOPN; ++s) {
            int bi = my_sel[s];
            if (bi < 0) continue;
            int s0 = bi * BS;
            if (s0 > t) continue;
            int s1 = min(s0 + BS, t + 1);
            for (int j = s0; j < s1; ++j) {
                PROCESS_KEY(j);
            }
        }

        // Pass B: window keys not in selected blocks
        for (int j = w0; j <= t; ++j) {
            if (is_sel(j / BS)) continue;
            PROCESS_KEY(j);
        }

        if (l_i == 0.f) {
            PROCESS_KEY(t);
        }
        #undef PROCESS_KEY

        const float inv = 1.f / l_i;
        #pragma unroll
        for (int e = 0; e < EPT; ++e) {
            int d = lane + e * 32;
            if (d < D_)
                o_base[(static_cast<long long>(t)) * D_ + d] = f2bf(acc[e] * inv);
        }
    }
}

// Generic fallback for uncommon D
__global__ void nsa_single_kernel(
    const __nv_bfloat16* __restrict__ Q,
    const __nv_bfloat16* __restrict__ K,
    const __nv_bfloat16* __restrict__ V,
    __nv_bfloat16* __restrict__ O,
    int B_, int H_, int S_, int D_
) {
    const int bh = blockIdx.x, t = blockIdx.y;
    if (t >= S_ || bh >= B_ * H_) return;
    const int tid = threadIdx.x, NT = blockDim.x;
    const long long head_off = (static_cast<long long>(bh) * S_) * D_;
    const __nv_bfloat16* q_ptr  = Q + head_off + (static_cast<long long>(t)) * D_;
    const __nv_bfloat16* k_base = K + head_off;
    const __nv_bfloat16* v_base = V + head_off;
    __nv_bfloat16* o_ptr = O + head_off + (static_cast<long long>(t)) * D_;
    const int n_blocks = (S_ + BS - 1) / BS;
    const float scale = rsqrtf(static_cast<float>(D_));
    extern __shared__ float smem[];
    float* q_f = smem;
    float* block_imp = q_f + D_;
    float* red = block_imp + n_blocks;
    int* top_idx = reinterpret_cast<int*>(red + 32);
    unsigned char* sel_mask = reinterpret_cast<unsigned char*>(top_idx + TOPN);
    for (int d = tid; d < D_; d += NT) q_f[d] = bf2f(q_ptr[d]);
    __syncthreads();
    auto bsum = [&](float v) -> float {
        int lane = tid & 31, wid = tid >> 5;
        v = warp_sum(v);
        if (lane == 0) red[wid] = v;
        __syncthreads();
        int nwarps = (NT + 31) >> 5;
        v = (tid < nwarps) ? red[tid] : 0.f;
        if (wid == 0) v = warp_sum(v);
        if (tid == 0) red[0] = v;
        __syncthreads();
        v = red[0];
        __syncthreads();
        return v;
    };
    for (int bi = 0; bi < n_blocks; ++bi) {
        int s0 = bi * BS;
        if (s0 > t) {
            if (tid == 0) block_imp[bi] = -1e9f;
            __syncthreads();
            continue;
        }
        int s1 = min(s0 + BS, t + 1);
        float local = 0.f;
        for (int j = s0 + tid; j < s1; j += NT) {
            const __nv_bfloat16* kj = k_base + (static_cast<long long>(j)) * D_;
            float dot = 0.f;
            for (int d = 0; d < D_; ++d) dot += q_f[d] * bf2f(kj[d]);
            local += dot;
        }
        float total = bsum(local);
        if (tid == 0) block_imp[bi] = (total / float(s1 - s0)) * scale;
        __syncthreads();
    }
    if (tid == 0) {
        int n_causal = 0;
        for (int bi = 0; bi < n_blocks; ++bi)
            if (bi * BS <= t) n_causal++; else break;
        int k_sel = TOPN < n_causal ? TOPN : n_causal;
        for (int s = 0; s < k_sel; ++s) {
            float best = -1e30f; int best_i = 0;
            for (int bi = 0; bi < n_causal; ++bi) {
                bool used = false;
                for (int p = 0; p < s; ++p) if (top_idx[p] == bi) used = true;
                if (!used && block_imp[bi] > best) { best = block_imp[bi]; best_i = bi; }
            }
            top_idx[s] = best_i;
        }
        for (int s = k_sel; s < TOPN; ++s) top_idx[s] = -1;
        for (int bi = 0; bi < n_blocks; ++bi) sel_mask[bi] = 0;
        for (int s = 0; s < k_sel; ++s) sel_mask[top_idx[s]] = 1;
    }
    __syncthreads();
    float acc[8]; for (int i = 0; i < 8; ++i) acc[i] = 0.f;
    float m_i = -FLT_MAX, l_i = 0.f;
    int w0 = (t + 1 > WIN) ? (t + 1 - WIN) : 0;
    float q_reg[8];
    for (int i = 0; i < 8; ++i) {
        int d = tid + i * NT;
        q_reg[i] = (d < D_) ? q_f[d] : 0.f;
    }
    for (int bi = 0; bi < n_blocks; ++bi) {
        int s0 = bi * BS;
        if (s0 > t) break;
        int s1 = min(s0 + BS, t + 1);
        int j0, j1;
        if (sel_mask[bi]) { j0 = s0; j1 = s1; }
        else {
            j0 = s0 > w0 ? s0 : w0; j1 = s1;
            if (j0 >= j1) continue;
        }
        for (int j = j0; j < j1; ++j) {
            const __nv_bfloat16* kj = k_base + (static_cast<long long>(j)) * D_;
            const __nv_bfloat16* vj = v_base + (static_cast<long long>(j)) * D_;
            float partial = 0.f;
            for (int i = 0; i < 8; ++i) {
                int d = tid + i * NT;
                if (d < D_) partial += q_reg[i] * bf2f(kj[d]);
            }
            float score = bsum(partial) * scale;
            float m_new = fmaxf(m_i, score);
            float ed = __expf(m_i - m_new);
            float es = __expf(score - m_new);
            for (int i = 0; i < 8; ++i) acc[i] *= ed;
            l_i = l_i * ed + es;
            m_i = m_new;
            for (int i = 0; i < 8; ++i) {
                int d = tid + i * NT;
                if (d < D_) acc[i] += es * bf2f(vj[d]);
            }
        }
    }
    if (l_i == 0.f) {
        const __nv_bfloat16* kj = k_base + (static_cast<long long>(t)) * D_;
        const __nv_bfloat16* vj = v_base + (static_cast<long long>(t)) * D_;
        float partial = 0.f;
        for (int i = 0; i < 8; ++i) {
            int d = tid + i * NT;
            if (d < D_) partial += q_reg[i] * bf2f(kj[d]);
        }
        float score = bsum(partial) * scale;
        float m_new = fmaxf(m_i, score);
        float ed = __expf(m_i - m_new);
        float es = __expf(score - m_new);
        for (int i = 0; i < 8; ++i) acc[i] *= ed;
        l_i = l_i * ed + es;
        for (int i = 0; i < 8; ++i) {
            int d = tid + i * NT;
            if (d < D_) acc[i] += es * bf2f(vj[d]);
        }
    }
    float inv = 1.f / l_i;
    for (int i = 0; i < 8; ++i) {
        int d = tid + i * NT;
        if (d < D_) o_ptr[d] = f2bf(acc[i] * inv);
    }
}

torch::Tensor nsa_forward(
    torch::Tensor q, torch::Tensor k, torch::Tensor v,
    int64_t, int64_t, int64_t
) {
    TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda(), "CUDA required");
    TORCH_CHECK(q.scalar_type() == at::kBFloat16, "bf16 required");
    TORCH_CHECK(q.dim() == 4, "(B,H,S,D)");

    auto qc = q.contiguous();
    auto kc = k.contiguous();
    auto vc = v.contiguous();
    const int B_ = qc.size(0), H_ = qc.size(1), S_ = qc.size(2), D_ = qc.size(3);
    auto out = torch::empty_like(qc);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    auto* Qp = reinterpret_cast<const __nv_bfloat16*>(qc.data_ptr<at::BFloat16>());
    auto* Kp = reinterpret_cast<const __nv_bfloat16*>(kc.data_ptr<at::BFloat16>());
    auto* Vp = reinterpret_cast<const __nv_bfloat16*>(vc.data_ptr<at::BFloat16>());
    auto* Op = reinterpret_cast<__nv_bfloat16*>(out.data_ptr<at::BFloat16>());

    auto launch = [&](auto kernel, int TD, int QT, int threads) {
        dim3 grid(B_ * H_, (S_ + QT - 1) / QT);
        size_t smem =
            sizeof(float) * ((size_t)QT * TD + (size_t)BS * TD + (size_t)QT) +
            sizeof(int) * (size_t)QT * TOPN;
        smem = (smem + 15) & ~size_t(15);
        cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem);
        kernel<<<grid, threads, smem, stream>>>(Qp, Kp, Vp, Op, B_, H_, S_, D_);
    };

    if (D_ == 64) {
        // QT=32, 256 threads: 8 warps → each warp handles 4 queries in phase 2
        // smem: q 8KB + k 16KB + sel ~1KB ≈ 25KB
        launch(nsa_kernel<64, 32, 256>, 64, 32, 256);
    } else if (D_ == 128) {
        // QT=16, 256 threads: more phase-2 parallelism
        // smem: q 8KB + k 32KB + sel ≈ 41KB
        launch(nsa_kernel<128, 16, 256>, 128, 16, 256);
    } else {
        dim3 grid(B_ * H_, S_);
        int threads = 128;
        int n_blocks = (S_ + BS - 1) / BS;
        size_t smem = sizeof(float) * (D_ + n_blocks + 32) + sizeof(int) * TOPN + n_blocks;
        smem = (smem + 15) & ~size_t(15);
        nsa_single_kernel<<<grid, threads, smem, stream>>>(Qp, Kp, Vp, Op, B_, H_, S_, D_);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

CPP_SRC = r"""
torch::Tensor nsa_forward(
    torch::Tensor q, torch::Tensor k, torch::Tensor v,
    int64_t block_size, int64_t top_n_blocks, int64_t sliding_window);
"""

_nsa_mod: Optional[object] = None


def _get_mod():
    global _nsa_mod
    if _nsa_mod is None:
        _nsa_mod = load_inline(
            name="nsa_sparse_attn_v7",
            cpp_sources=[CPP_SRC],
            cuda_sources=[CUDA_SRC],
            functions=["nsa_forward"],
            extra_cuda_cflags=[
                "-O3",
                "--use_fast_math",
                "-U__CUDA_NO_HALF_OPERATORS__",
                "-U__CUDA_NO_HALF_CONVERSIONS__",
                "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
                "-U__CUDA_NO_HALF2_OPERATORS__",
                "--expt-relaxed-constexpr",
            ],
            verbose=False,
        )
    return _nsa_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))
        if torch.cuda.is_available():
            _get_mod()

    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
        return _get_mod().nsa_forward(
            q, k, v, BLOCK_SIZE, TOP_N_BLOCKS, SLIDING_WINDOW
        )


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]