"""High-performance top-k via custom CUDA adaptive radix-select.

Strategies:
  k=1     — vectorized block argmax
  k=8, n≤8192 — single-pass per-thread insert + tree merge
  else    — adaptive MSB radix-select with early materialize + bitonic finish
"""
from __future__ import annotations

from pathlib import Path

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

CUDA_SRC = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cfloat>
#include <cstdint>

struct Pair {
    float val;
    int   idx;
};

__device__ __forceinline__ bool pgt(Pair a, Pair b) {
    return (a.val > b.val) || (a.val == b.val && a.idx < b.idx);
}

// Map float to uint so unsigned order matches float order.
__device__ __forceinline__ uint32_t f2o(float f) {
    uint32_t b = __float_as_uint(f);
    return b ^ ((b >> 31) ? 0xFFFFFFFFu : 0x80000000u);
}

// =====================================================================
// k = 1 specialized argmax
// =====================================================================
__global__ void k1_kernel(
    const float* __restrict__ x, float* __restrict__ ov, int64_t* __restrict__ oi,
    int n, int batch)
{
    int row = blockIdx.x;
    if (row >= batch) return;
    const float* r = x + (size_t)row * n;
    float bv = -FLT_MAX; int bi = 0;
    int n4 = n >> 2;
    const float4* r4 = reinterpret_cast<const float4*>(r);
    for (int i = threadIdx.x; i < n4; i += blockDim.x) {
        float4 v = __ldg(r4 + i);
        int b = i << 2;
        if (v.x > bv) { bv = v.x; bi = b; }
        if (v.y > bv) { bv = v.y; bi = b + 1; }
        if (v.z > bv) { bv = v.z; bi = b + 2; }
        if (v.w > bv) { bv = v.w; bi = b + 3; }
    }
    for (int i = (n4 << 2) + threadIdx.x; i < n; i += blockDim.x) {
        float v = __ldg(r + i);
        if (v > bv) { bv = v; bi = i; }
    }
    unsigned m = 0xffffffffu;
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) {
        float ovv = __shfl_down_sync(m, bv, o);
        int oi2 = __shfl_down_sync(m, bi, o);
        if (ovv > bv || (ovv == bv && oi2 < bi)) { bv = ovv; bi = oi2; }
    }
    __shared__ float sv[32]; __shared__ int si[32];
    int lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
    int nw = (blockDim.x + 31) >> 5;
    if (lane == 0) { sv[warp] = bv; si[warp] = bi; }
    __syncthreads();
    if (warp == 0) {
        bv = (lane < nw) ? sv[lane] : -FLT_MAX;
        bi = (lane < nw) ? si[lane] : 0;
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) {
            float ovv = __shfl_down_sync(m, bv, o);
            int oi2 = __shfl_down_sync(m, bi, o);
            if (ovv > bv || (ovv == bv && oi2 < bi)) { bv = ovv; bi = oi2; }
        }
        if (lane == 0) { ov[row] = bv; oi[row] = (int64_t)bi; }
    }
}

// =====================================================================
// Bitonic sort ascending on shared Pair array
// =====================================================================
template<int BLOCK>
__device__ void bitonic_asc(Pair* a, int pow2) {
    for (int size = 2; size <= pow2; size <<= 1) {
        for (int stride = size >> 1; stride > 0; stride >>= 1) {
            for (int i = threadIdx.x; i < pow2; i += BLOCK) {
                int j = i ^ stride;
                if (j > i) {
                    bool asc = ((i & size) == 0);
                    if (asc) {
                        if (pgt(a[i], a[j])) { Pair t = a[i]; a[i] = a[j]; a[j] = t; }
                    } else {
                        if (pgt(a[j], a[i])) { Pair t = a[i]; a[i] = a[j]; a[j] = t; }
                    }
                }
            }
            __syncthreads();
        }
    }
}

// =====================================================================
// Adaptive radix top-K — one block per row
// =====================================================================
template<int BLOCK>
__global__ void adaptive_radix_topk(
    const float* __restrict__ x,
    float* __restrict__ ov, int64_t* __restrict__ oi,
    int n, int batch, int K)
{
    constexpr int NW = BLOCK / 32;
    constexpr int MAXC = 512;
    constexpr int THRESH = 320;

    __shared__ int whist[NW][256];
    __shared__ int hist[256];
    __shared__ Pair cands[MAXC];
    __shared__ int s_count, s_bucket, s_kbits, s_rank;
    __shared__ uint32_t s_known;

    int row = blockIdx.x;
    if (row >= batch) return;
    const float* r = x + (size_t)row * n;
    int lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
    int n4 = n >> 2;
    const float4* r4 = reinterpret_cast<const float4*>(r);

    if (threadIdx.x == 0) { s_known = 0; s_kbits = 0; s_rank = K; }
    __syncthreads();

    for (int pass = 0; pass < 4; ++pass) {
        int shift = 24 - 8 * pass;
        int kbits = s_kbits;
        uint32_t known = s_known;

        for (int b = lane; b < 256; b += 32) whist[warp][b] = 0;
        __syncthreads();

        for (int i = threadIdx.x; i < n4; i += BLOCK) {
            float4 v = __ldg(r4 + i);
            float vals[4] = {v.x, v.y, v.z, v.w};
            #pragma unroll
            for (int j = 0; j < 4; ++j) {
                uint32_t key = f2o(vals[j]);
                bool match = (kbits == 0) ||
                    ((key >> (32 - kbits)) == (known >> (32 - kbits)));
                if (match) atomicAdd(&whist[warp][(key >> shift) & 0xFFu], 1);
            }
        }
        for (int i = (n4 << 2) + threadIdx.x; i < n; i += BLOCK) {
            uint32_t key = f2o(__ldg(r + i));
            bool match = (kbits == 0) ||
                ((key >> (32 - kbits)) == (known >> (32 - kbits)));
            if (match) atomicAdd(&whist[warp][(key >> shift) & 0xFFu], 1);
        }
        __syncthreads();

        for (int b = threadIdx.x; b < 256; b += BLOCK) {
            int s = 0;
            #pragma unroll
            for (int w = 0; w < NW; ++w) s += whist[w][b];
            hist[b] = s;
        }
        __syncthreads();

        if (threadIdx.x == 0) {
            int cum = 0, dig = 0, ab = 0, rank = s_rank;
            for (dig = 255; dig >= 0; --dig) {
                int c = hist[dig];
                if (cum + c >= rank) { ab = cum; break; }
                cum += c;
            }
            s_bucket = hist[dig];
            s_rank = rank - ab;
            s_known = known | ((uint32_t)dig << shift);
            s_kbits = kbits + 8;
        }
        __syncthreads();

        if ((s_bucket <= THRESH) || (pass == 3)) {
            if (threadIdx.x == 0) s_count = 0;
            __syncthreads();
            int kbits2 = s_kbits;
            uint32_t pref = (kbits2 == 0) ? 0u : (s_known >> (32 - kbits2));

            for (int i = threadIdx.x; i < n4; i += BLOCK) {
                float4 v = __ldg(r4 + i);
                float vals[4] = {v.x, v.y, v.z, v.w};
                #pragma unroll
                for (int j = 0; j < 4; ++j) {
                    uint32_t key = f2o(vals[j]);
                    uint32_t kp = (kbits2 == 0) ? 0u : (key >> (32 - kbits2));
                    if (kp >= pref) {
                        int slot = atomicAdd(&s_count, 1);
                        if (slot < MAXC) cands[slot] = Pair{vals[j], (i << 2) + j};
                    }
                }
            }
            for (int i = (n4 << 2) + threadIdx.x; i < n; i += BLOCK) {
                float v = __ldg(r + i);
                uint32_t key = f2o(v);
                uint32_t kp = (kbits2 == 0) ? 0u : (key >> (32 - kbits2));
                if (kp >= pref) {
                    int slot = atomicAdd(&s_count, 1);
                    if (slot < MAXC) cands[slot] = Pair{v, i};
                }
            }
            __syncthreads();

            if (s_count > MAXC && pass < 3) {
                // Bucket still too big — refine more digits.
            } else {
                int nc = s_count; if (nc > MAXC) nc = MAXC;
                int need = nc > K ? nc : K;
                int pow2 = 1; while (pow2 < need) pow2 <<= 1;
                if (pow2 > MAXC) pow2 = MAXC;
                for (int i = threadIdx.x + nc; i < pow2; i += BLOCK)
                    cands[i] = Pair{-FLT_MAX, 0x7fffffff};
                __syncthreads();
                bitonic_asc<BLOCK>(cands, pow2);
                if (threadIdx.x < K) {
                    Pair p = cands[pow2 - 1 - threadIdx.x];
                    ov[row * K + threadIdx.x] = p.val;
                    oi[row * K + threadIdx.x] = (int64_t)p.idx;
                }
                return;
            }
        }
        __syncthreads();
    }
}

// =====================================================================
// Small-K single-pass select (K=8)
// =====================================================================
template<int K>
__device__ __forceinline__ void insertK(Pair* a, Pair p) {
    if (p.val < a[K - 1].val) return;
    if (p.val == a[K - 1].val && p.idx >= a[K - 1].idx) return;
    #pragma unroll
    for (int j = K - 1; j > 0; --j) {
        if (pgt(p, a[j - 1])) a[j] = a[j - 1];
        else { a[j] = p; return; }
    }
    a[0] = p;
}

template<int K>
__device__ __forceinline__ void mergeK(Pair* a, const Pair* b, Pair* t) {
    int ia = 0, ib = 0;
    #pragma unroll
    for (int i = 0; i < K; ++i) {
        if (ia < K && (ib >= K || pgt(a[ia], b[ib]))) t[i] = a[ia++];
        else t[i] = b[ib++];
    }
    #pragma unroll
    for (int i = 0; i < K; ++i) a[i] = t[i];
}

template<int K, int BLOCK>
__global__ void select_kernel(
    const float* __restrict__ x, float* __restrict__ ov, int64_t* __restrict__ oi,
    int n, int batch)
{
    extern __shared__ char raw[];
    Pair* sm = reinterpret_cast<Pair*>(raw);
    int row = blockIdx.x;
    if (row >= batch) return;
    const float* r = x + (size_t)row * n;

    Pair local[K];
    #pragma unroll
    for (int j = 0; j < K; ++j) local[j] = Pair{-FLT_MAX, 0x7fffffff};

    int n4 = n >> 2;
    const float4* r4 = reinterpret_cast<const float4*>(r);
    for (int i = threadIdx.x; i < n4; i += BLOCK) {
        float4 v = __ldg(r4 + i);
        int b = i << 2;
        insertK<K>(local, Pair{v.x, b});
        insertK<K>(local, Pair{v.y, b + 1});
        insertK<K>(local, Pair{v.z, b + 2});
        insertK<K>(local, Pair{v.w, b + 3});
    }
    for (int i = (n4 << 2) + threadIdx.x; i < n; i += BLOCK)
        insertK<K>(local, Pair{__ldg(r + i), i});

    #pragma unroll
    for (int j = 0; j < K; ++j) sm[threadIdx.x * K + j] = local[j];
    __syncthreads();

    Pair tmp[K];
    for (int s = BLOCK >> 1; s >= 1; s >>= 1) {
        if (threadIdx.x < s)
            mergeK<K>(&sm[threadIdx.x * K], &sm[(threadIdx.x + s) * K], tmp);
        __syncthreads();
    }
    for (int j = threadIdx.x; j < K; j += BLOCK) {
        ov[row * K + j] = sm[j].val;
        oi[row * K + j] = (int64_t)sm[j].idx;
    }
}

// =====================================================================
// Host dispatcher
// =====================================================================
void topk_cuda(
    torch::Tensor x, torch::Tensor values, torch::Tensor indices, int64_t k)
{
    TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kFloat32);
    TORCH_CHECK(x.dim() == 2);

    const int batch = (int)x.size(0);
    const int n     = (int)x.size(1);
    const int K     = (int)k;
    const float* xp = x.data_ptr<float>();
    float* vp = values.data_ptr<float>();
    int64_t* ip = indices.data_ptr<int64_t>();

    if (K == 1) {
        k1_kernel<<<batch, (n >= 2048) ? 512 : 256>>>(xp, vp, ip, n, batch);
        return;
    }

    // Single-pass insert is very competitive for tiny K and moderate n
    if (K == 8 && n <= 8192) {
        size_t smem = 256ull * 8 * sizeof(Pair);
        select_kernel<8, 256> <<<batch, 256, smem>>>(xp, vp, ip, n, batch);
        return;
    }

    if (n >= 4096)
        adaptive_radix_topk<512> <<<batch, 512>>>(xp, vp, ip, n, batch, K);
    else
        adaptive_radix_topk<256> <<<batch, 256>>>(xp, vp, ip, n, batch, K);
}
"""

_ext = None


def _get_ext():
    global _ext
    if _ext is not None:
        return _ext
    build_dir = Path(__file__).resolve().parent / ".topk_ext_build"
    build_dir.mkdir(exist_ok=True)
    major, minor = torch.cuda.get_device_capability(0)
    arch = f"{major}{minor}"
    _ext = load_inline(
        name="topk_bitonic_ext",
        cpp_sources=[
            "void topk_cuda(torch::Tensor x, torch::Tensor values, "
            "torch::Tensor indices, int64_t k);"
        ],
        cuda_sources=[CUDA_SRC],
        functions=["topk_cuda"],
        extra_cuda_cflags=[
            "-O3",
            "--use_fast_math",
            f"-gencode=arch=compute_{arch},code=sm_{arch}",
        ],
        extra_cflags=["-O3"],
        build_directory=str(build_dir),
        verbose=False,
    )
    return _ext


class Model(nn.Module):
    def __init__(self, batch: int, n: int, k: int):
        super().__init__()
        self.batch, self.n, self.k = batch, n, k
        self.register_buffer("_dummy", torch.zeros(1))
        if torch.cuda.is_available():
            _get_ext()

    def forward(self, x: torch.Tensor):
        B, _ = x.shape
        K = self.k
        values = torch.empty(B, K, device=x.device, dtype=torch.float32)
        indices = torch.empty(B, K, device=x.device, dtype=torch.int64)
        _get_ext().topk_cuda(x.contiguous(), values, indices, int(K))
        return values, indices


batch = 64
n = 8192
k = 8


def get_inputs():
    return [torch.randn(batch, n, dtype=torch.float32)]


def get_init_inputs():
    return [batch, n, k]