"""Top-k via bitonic sort in shared memory, targeting B200 (SM100, HBM3e).

Strategy:
  - Rows that fit in shared memory (n <= 16384): one CUDA block per row, bitonic
    sort on padded-to-power-of-2 (value, index) pairs in shared memory, then
    read off the first k elements.
  - Large rows (n = 131072): split into 8 segments of 16384, sort each in
    parallel, then merge the partial top-k lists (8*k candidates) with a tiny
    second kernel.
  - k=1 (argmax): dedicated reduction kernel that avoids the full sort.
"""

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

# ---------------------------------------------------------------------------
# CUDA source
# ---------------------------------------------------------------------------

CUDA_SOURCE = r"""
#include <cuda_runtime.h>
#include <math.h>

// Negative infinity for fp32; works in device code.
#define NEG_INF (-1.0f / 0.0f)

struct Pair {
    float value;
    int   index;
};

// ------------------------------------------------------------------
// Device helper: bitonic sort.  N_POW2 and BLOCK_DIM are template
// params so the compiler can unroll the loop nests.
// ------------------------------------------------------------------
template <int N_POW2, int BLOCK_DIM>
__device__ void bitonic_sort(Pair* shmem) {
    #pragma unroll
    for (int k = 2; k <= N_POW2; k <<= 1) {
        #pragma unroll
        for (int j = k >> 1; j > 0; j >>= 1) {
            for (int tid = threadIdx.x; tid < (N_POW2 >> 1); tid += BLOCK_DIM) {
                int ixj = tid ^ j;
                if (ixj > tid) {
                    float a = shmem[tid].value;
                    float b = shmem[ixj].value;
                    bool do_swap;
                    if ((tid & k) == 0) {
                        do_swap = (a < b);   // descending
                    } else {
                        do_swap = (a > b);   // ascending
                    }
                    if (do_swap) {
                        Pair t = shmem[tid];
                        shmem[tid]      = shmem[ixj];
                        shmem[ixj]      = t;
                    }
                }
            }
            __syncthreads();
        }
    }
}

// ------------------------------------------------------------------
// Explicitly-instantiated kernels  (one per N_POW2 we use)
// ------------------------------------------------------------------

// ---- 4096 ----
extern "C" __global__ void bitonic_topk_4096(
    const float* __restrict__ input,
    float*       __restrict__ out_values,
    long long*   __restrict__ out_indices,
    int n, int k, int stride_n)
{
    constexpr int N = 4096, B = 256;
    extern __shared__ char smem[];
    Pair* shmem = (Pair*)smem;
    int bid = blockIdx.x;
    const float* row = input + bid * stride_n;
    int chunk = N / B;
    for (int i = threadIdx.x * chunk, end = i + chunk; i < end; ++i) {
        shmem[i] = {i < n ? row[i] : NEG_INF, i < n ? i : -1};
    }
    __syncthreads();
    bitonic_sort<N, B>(shmem);
    __syncthreads();
    int off = bid * k;
    for (int i = threadIdx.x; i < k; i += B)
        out_values[off+i] = shmem[i].value, out_indices[off+i] = (long long)shmem[i].index;
}

// ---- 8192 ----
extern "C" __global__ void bitonic_topk_8192(
    const float* __restrict__ input,
    float*       __restrict__ out_values,
    long long*   __restrict__ out_indices,
    int n, int k, int stride_n)
{
    constexpr int N = 8192, B = 512;
    extern __shared__ char smem[];
    Pair* shmem = (Pair*)smem;
    int bid = blockIdx.x;
    const float* row = input + bid * stride_n;
    int chunk = N / B;
    for (int i = threadIdx.x * chunk, end = i + chunk; i < end; ++i) {
        shmem[i] = {i < n ? row[i] : NEG_INF, i < n ? i : -1};
    }
    __syncthreads();
    bitonic_sort<N, B>(shmem);
    __syncthreads();
    int off = bid * k;
    for (int i = threadIdx.x; i < k; i += B)
        out_values[off+i] = shmem[i].value, out_indices[off+i] = (long long)shmem[i].index;
}

// ---- 16384 ----
extern "C" __global__ void bitonic_topk_16384(
    const float* __restrict__ input,
    float*       __restrict__ out_values,
    long long*   __restrict__ out_indices,
    int n, int k, int stride_n)
{
    constexpr int N = 16384, B = 512;
    extern __shared__ char smem[];
    Pair* shmem = (Pair*)smem;
    int bid = blockIdx.x;
    const float* row = input + bid * stride_n;
    int chunk = N / B;
    for (int i = threadIdx.x * chunk, end = i + chunk; i < end; ++i) {
        shmem[i] = {i < n ? row[i] : NEG_INF, i < n ? i : -1};
    }
    __syncthreads();
    bitonic_sort<N, B>(shmem);
    __syncthreads();
    int off = bid * k;
    for (int i = threadIdx.x; i < k; i += B)
        out_values[off+i] = shmem[i].value, out_indices[off+i] = (long long)shmem[i].index;
}

// ---- 16384 segmented (multiple blocks share one row) ----
extern "C" __global__ void bitonic_seg_16384(
    const float* __restrict__ input,
    float*       __restrict__ partial_vals,
    long long*   __restrict__ partial_idx,
    int n, int k, int stride_n, int num_segs)
{
    constexpr int S = 16384, B = 512;
    extern __shared__ char smem[];
    Pair* shmem = (Pair*)smem;
    int bid = blockIdx.y;
    int seg = blockIdx.x;
    int off = seg * S;
    int seg_n = off < n ? min(S, n - off) : 0;
    const float* row = input + bid * stride_n + off;
    int chunk = S / B;
    for (int i = threadIdx.x * chunk, end = i + chunk; i < end; ++i) {
        shmem[i] = {i < seg_n ? row[i] : NEG_INF, i < seg_n ? i : -1};
    }
    __syncthreads();
    bitonic_sort<S, B>(shmem);
    __syncthreads();
    int out_off = (bid * num_segs + seg) * k;
    for (int i = threadIdx.x; i < k; i += B) {
        partial_vals[out_off+i] = shmem[i].value;
        int idx = shmem[i].index;
        partial_idx[out_off+i] = (idx >= 0) ? (long long)(off + idx) : (long long)-1;
    }
}

// ---- Merge partial results ----
extern "C" __global__ void merge_topk(
    const float*    __restrict__ partial_vals,
    const long long*__restrict__ partial_idx,
    float*          __restrict__ out_values,
    long long*      __restrict__ out_indices,
    int total_candidates, int k)
{
    extern __shared__ char smem[];
    Pair* shmem = (Pair*)smem;
    int bid = blockIdx.x;
    int src_off = bid * total_candidates;

    int N = 1; while (N < total_candidates) N <<= 1;

    for (int i = threadIdx.x; i < N; i += blockDim.x)
        shmem[i] = {i < total_candidates ? partial_vals[src_off+i] : NEG_INF,
                    i < total_candidates ? (int)partial_idx[src_off+i] : -1};
    __syncthreads();

    for (int kk = 2; kk <= N; kk <<= 1) {
        for (int j = kk >> 1; j > 0; j >>= 1) {
            for (int tid = threadIdx.x; tid < (N >> 1); tid += blockDim.x) {
                int ixj = tid ^ j;
                if (ixj <= tid) continue;
                float a = shmem[tid].value, b = shmem[ixj].value;
                bool do_swap = (tid & kk) == 0 ? (a < b) : (a > b);
                if (do_swap) { Pair t = shmem[tid]; shmem[tid] = shmem[ixj]; shmem[ixj] = t; }
            }
            __syncthreads();
        }
    }

    for (int i = threadIdx.x; i < k; i += blockDim.x) {
        out_values[bid * k + i] = shmem[i].value;
        out_indices[bid * k + i] = (long long)shmem[i].index;
    }
}

// ---- Argmax fast path ----
extern "C" __global__ void argmax_kernel(
    const float* __restrict__ input,
    float*       __restrict__ out_values,
    long long*   __restrict__ out_indices,
    int n, int stride_n)
{
    __shared__ Pair smem[512];
    int bid = blockIdx.x;
    const float* row = input + bid * stride_n;
    float best = NEG_INF; int best_i = -1;
    for (int i = threadIdx.x; i < n; i += blockDim.x) {
        float v = row[i];
        if (v > best) { best = v; best_i = i; }
    }
    smem[threadIdx.x] = {best, best_i};
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (threadIdx.x < s && smem[threadIdx.x+s].value > smem[threadIdx.x].value)
            smem[threadIdx.x] = smem[threadIdx.x+s];
        __syncthreads();
    }
    if (threadIdx.x == 0) {
        out_values[bid] = smem[0].value;
        out_indices[bid] = (long long)smem[0].index;
    }
}
"""

# ---------------------------------------------------------------------------
# Compile once
# ---------------------------------------------------------------------------
_topk_ops = None


def _get_ops():
    global _topk_ops
    if _topk_ops is not None:
        return _topk_ops
    _topk_ops = load_inline(
        name="topk_bitonic",
        cpp_sources="",
        cuda_sources=CUDA_SOURCE,
        functions=[
            "bitonic_topk_4096",
            "bitonic_topk_8192",
            "bitonic_topk_16384",
            "bitonic_seg_16384",
            "merge_topk",
            "argmax_kernel",
        ],
        extra_cuda_cflags=["-arch=sm_100", "--expt-relaxed-constexpr"],
        verbose=True,
    )
    return _topk_ops


# ---------------------------------------------------------------------------
# Python dispatch
# ---------------------------------------------------------------------------

def _next_pow2(x: int) -> int:
    p = 1
    while p < x:
        p <<= 1
    return p


def _compute_topk(x: torch.Tensor, k: int) -> tuple[torch.Tensor, torch.Tensor]:
    ops = _get_ops()
    batch, n = x.shape
    n_pow2 = _next_pow2(n)
    stride = x.stride(0)
    device = x.device
    stream = torch.cuda.current_stream(device)

    values = torch.empty(batch, k, dtype=torch.float32, device=device)
    indices = torch.empty(batch, k, dtype=torch.int64, device=device)

    # --- Argmax fast path ---
    if k == 1 and n_pow2 <= 16384:
        ops.argmax_kernel(
            (batch,), (512,), (512 * 8,), stream,
            x.data_ptr(), values.data_ptr(), indices.data_ptr(),
            n, stride,
        )
        return values, indices

    # --- Single-block per row ---
    if n_pow2 <= 16384:
        if n_pow2 == 4096:
            ops.bitonic_topk_4096(
                (batch,), (256,), (4096 * 8,), stream,
                x.data_ptr(), values.data_ptr(), indices.data_ptr(),
                n, k, stride,
            )
        elif n_pow2 == 8192:
            ops.bitonic_topk_8192(
                (batch,), (512,), (8192 * 8,), stream,
                x.data_ptr(), values.data_ptr(), indices.data_ptr(),
                n, k, stride,
            )
        else:
            ops.bitonic_topk_16384(
                (batch,), (512,), (16384 * 8,), stream,
                x.data_ptr(), values.data_ptr(), indices.data_ptr(),
                n, k, stride,
            )
        return values, indices

    # --- Segmented: split n into chunks of 16384 ---
    seg_size = 16384
    num_segs = (n + seg_size - 1) // seg_size

    partial_vals = torch.empty(batch, num_segs * k, dtype=torch.float32, device=device)
    partial_idx = torch.empty(batch, num_segs * k, dtype=torch.int64, device=device)

    ops.bitonic_seg_16384(
        (num_segs, batch), (512,), (16384 * 8,), stream,
        x.data_ptr(), partial_vals.data_ptr(), partial_idx.data_ptr(),
        n, k, stride, num_segs,
    )

    total = num_segs * k
    merge_pow2 = _next_pow2(total)
    merge_block = min(merge_pow2 // 2, 256)
    if merge_block < 32:
        merge_block = 32

    ops.merge_topk(
        (batch,), (merge_block,), (merge_pow2 * 8,), stream,
        partial_vals.data_ptr(), partial_idx.data_ptr(),
        values.data_ptr(), indices.data_ptr(),
        total, k,
    )

    return values, indices


# ---------------------------------------------------------------------------
# Model  (same interface as reference.py)
# ---------------------------------------------------------------------------

OP_TYPE = "topk"
SUPPORTED_PRECISIONS = ["fp32"]
HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]


class Model(nn.Module):
    """Top-k via bitonic sort in shared memory."""

    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))
        _get_ops()  # ensure compiled on first construction

    def forward(self, x: torch.Tensor):
        return _compute_topk(x, self.k)


# Module-level shims (rebuilt per shape by check.py / benchmark.py)
batch = 64
n = 8192
k = 8


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


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