"""Top-k via CUDA with CUB radix sort for segments and bitonic merge.

Uses CUB's highly-optimized BlockRadixSort for segment sorting (~20 syncs
vs 55 for hand-rolled bitonic), then a merge tree for combining segments.
"""

from __future__ import annotations

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

# ---------------------------------------------------------------------------
# CUDA source — uses CUB BlockRadixSort for fast segment sorting
# ---------------------------------------------------------------------------
CUDA_SRC = r"""
#include <cuda_runtime.h>
#include <cstdint>
#include <float.h>
#include <cub/block/block_radix_sort.cuh>

#define SEG_SIZE 1024
#define MAX_K 64
#define MERGE_BS 1024

// =========================================================================
// segment_kernel — uses CUB BlockRadixSort
//   grid = (batch * num_seg)   block = (SEG_SIZE)
//   Sorts one segment per block descending by value, writes top-k.
// =========================================================================
extern "C" __global__ void segment_kernel(
    const float* __restrict__ x,
    float* __restrict__ seg_vals,
    long long* __restrict__ seg_idxs,
    int n, int k, int num_seg)
{
    int row = blockIdx.x / num_seg;
    int seg = blockIdx.x % num_seg;
    int tid = threadIdx.x;

    // --- CUB BlockRadixSort ---
    typedef cub::BlockRadixSort<float, SEG_SIZE, 1, long long> BlockRadixSort;
    __shared__ union {
        typename BlockRadixSort::TempStorage sort;
        struct {
            float     vals[SEG_SIZE];
            long long idxs[SEG_SIZE];
        } data;
    } smem;

    // Load segment
    int start = seg * SEG_SIZE;
    int gidx = start + tid;
    float key;
    long long val;
    if (gidx < n) {
        key = x[row * n + gidx];
        val = gidx;
    } else {
        key = -INFINITY;
        val = -1;
    }

    // Radix sort descending by key (value).  After sort, largest keys
    // come first (descending order).
    // CUB 3.x API (CUDA 13) requires array params + begin_bit/end_bit.
    float keys_arr[1] = {key};
    long long vals_arr[1] = {val};
    BlockRadixSort(smem.sort).SortDescending(
        keys_arr, vals_arr, 0, 8 * sizeof(float));
    key = keys_arr[0];
    val = vals_arr[0];

    // Now key/val are in registers (striped layout), sorted descending.
    // Store to shared memory for top-k extraction.
    smem.data.vals[tid] = key;
    smem.data.idxs[tid] = val;
    __syncthreads();

    // Write top-k (already sorted descending by CUB, so first k are the largest)
    if (tid < k) {
        int dst = (row * num_seg + seg) * k + tid;
        seg_vals[dst] = smem.data.vals[tid];
        seg_idxs[dst] = smem.data.idxs[tid];
    }
}

// -------------------------------------------------------------------------
// bitonic_sort_asc — used in merge kernels (small N, cheap)
// -------------------------------------------------------------------------
__device__ void bitonic_sort_asc(
    float* __restrict__ vals,
    long long* __restrict__ idxs,
    int N)
{
    int tid = threadIdx.x;
    bool active = (tid < N);

    for (int stage = 2; stage <= N; stage <<= 1) {
        for (int step = stage >> 1; step > 0; step >>= 1) {
            int partner = tid ^ step;
            if (active && partner < N) {
                float vp = vals[partner];
                long long ip = idxs[partner];
                float vs = vals[tid];
                long long is_ = idxs[tid];

                bool ascending = (tid & stage) == 0;
                bool do_swap;

                if (ascending) {
                    do_swap = (vs > vp);
                    if (vs != vs) do_swap = true;
                    if (vp != vp) do_swap = false;
                } else {
                    do_swap = (vs < vp);
                    if (vs != vs) do_swap = true;
                    if (vp != vp) do_swap = false;
                }

                if (tid < partner) {
                    if (do_swap) {
                        vals[tid] = vp;
                        idxs[tid] = ip;
                        vals[partner] = vs;
                        idxs[partner] = is_;
                    }
                }
            }
            __syncthreads();
        }
    }
}

// =========================================================================
// merge_kernel
//   grid = (batch * num_out_seg)   block = (MERGE_BS=1024)
//   Merges group_size input segments (each with k elements, sorted desc)
//   into one output segment with k elements sorted ascending.
//
//   After merge + sort, writes top-k in ascending order.
// =========================================================================
extern "C" __global__ void merge_kernel(
    const float* __restrict__ in_vals,
    const long long* __restrict__ in_idxs,
    float* __restrict__ out_vals,
    long long* __restrict__ out_idxs,
    int k, int num_in_seg, int group_size, int num_out_seg)
{
    int row = blockIdx.x / num_out_seg;
    int out_seg = blockIdx.x % num_out_seg;
    int tid = threadIdx.x;

    __shared__ float     mvals[MERGE_BS];
    __shared__ long long midxs[MERGE_BS];

    int merge_N = group_size * k;  // elements to sort (<= MERGE_BS)

    // Load.  Input segments are sorted descending (from CUB).
    // We load them and will sort ascending.
    if (tid < merge_N) {
        int in_seg = out_seg * group_size + (tid / k);
        int in_k   = tid % k;
        if (in_seg < num_in_seg) {
            int src = (row * num_in_seg + in_seg) * k + in_k;
            mvals[tid] = in_vals[src];
            midxs[tid] = in_idxs[src];
        } else {
            mvals[tid] = -INFINITY;
            midxs[tid] = -1;
        }
    } else {
        mvals[tid] = -INFINITY;
        midxs[tid] = -1;
    }
    __syncthreads();

    // Sort ascending
    bitonic_sort_asc(mvals, midxs, MERGE_BS);

    // Write top-k (largest k = at end of ascending array)
    if (tid < k) {
        int src = MERGE_BS - k + tid;
        int dst = (row * num_out_seg + out_seg) * k + tid;
        out_vals[dst] = mvals[src];
        out_idxs[dst] = midxs[src];
    }
}

// =========================================================================
// final_write_kernel
//   grid = (batch)   block = (k)
//   Writes final result in descending order.
// =========================================================================
extern "C" __global__ void final_write_kernel(
    const float* __restrict__ in_vals,
    const long long* __restrict__ in_idxs,
    float* __restrict__ out_vals,
    long long* __restrict__ out_idxs,
    int k)
{
    int row = blockIdx.x;
    int tid = threadIdx.x;

    if (tid < k) {
        // Input is ascending (from merge). Write descending.
        int src = row * k + tid;
        int dst = row * k + (k - 1 - tid);
        out_vals[dst] = in_vals[src];
        out_idxs[dst] = in_idxs[src];
    }
}

// =========================================================================
// Launchers
// =========================================================================
extern "C" void launch_segment_kernel(
    unsigned long long x_ptr,
    unsigned long long seg_vals_ptr,
    unsigned long long seg_idxs_ptr,
    int batch, int n, int k, int num_seg,
    unsigned long long stream_ptr)
{
    cudaStream_t stream = reinterpret_cast<cudaStream_t>(stream_ptr);
    const float* x = reinterpret_cast<const float*>(x_ptr);
    float* seg_vals = reinterpret_cast<float*>(seg_vals_ptr);
    long long* seg_idxs = reinterpret_cast<long long*>(seg_idxs_ptr);

    dim3 grid(batch * num_seg);
    dim3 block(SEG_SIZE);
    segment_kernel<<<grid, block, 0, stream>>>(x, seg_vals, seg_idxs, n, k, num_seg);
}

extern "C" void launch_merge_kernel(
    unsigned long long in_vals_ptr,
    unsigned long long in_idxs_ptr,
    unsigned long long out_vals_ptr,
    unsigned long long out_idxs_ptr,
    int batch, int k, int num_in_seg, int group_size, int num_out_seg,
    unsigned long long stream_ptr)
{
    cudaStream_t stream = reinterpret_cast<cudaStream_t>(stream_ptr);
    const float* in_vals = reinterpret_cast<const float*>(in_vals_ptr);
    const long long* in_idxs = reinterpret_cast<const long long*>(in_idxs_ptr);
    float* out_vals = reinterpret_cast<float*>(out_vals_ptr);
    long long* out_idxs = reinterpret_cast<long long*>(out_idxs_ptr);

    dim3 grid(batch * num_out_seg);
    dim3 block(MERGE_BS);
    merge_kernel<<<grid, block, 0, stream>>>(
        in_vals, in_idxs, out_vals, out_idxs,
        k, num_in_seg, group_size, num_out_seg);
}

extern "C" void launch_final_write_kernel(
    unsigned long long in_vals_ptr,
    unsigned long long in_idxs_ptr,
    unsigned long long out_vals_ptr,
    unsigned long long out_idxs_ptr,
    int batch, int k,
    unsigned long long stream_ptr)
{
    cudaStream_t stream = reinterpret_cast<cudaStream_t>(stream_ptr);
    const float* in_vals = reinterpret_cast<const float*>(in_vals_ptr);
    const long long* in_idxs = reinterpret_cast<const long long*>(in_idxs_ptr);
    float* out_vals = reinterpret_cast<float*>(out_vals_ptr);
    long long* out_idxs = reinterpret_cast<long long*>(out_idxs_ptr);

    dim3 grid(batch);
    dim3 block(k);
    final_write_kernel<<<grid, block, 0, stream>>>(
        in_vals, in_idxs, out_vals, out_idxs, k);
}
"""

CPP_SRC = """
#include <cstdint>

extern "C" void launch_segment_kernel(
    unsigned long long x_ptr,
    unsigned long long seg_vals_ptr,
    unsigned long long seg_idxs_ptr,
    int batch, int n, int k, int num_seg,
    unsigned long long stream_ptr);

extern "C" void launch_merge_kernel(
    unsigned long long in_vals_ptr,
    unsigned long long in_idxs_ptr,
    unsigned long long out_vals_ptr,
    unsigned long long out_idxs_ptr,
    int batch, int k, int num_in_seg, int group_size, int num_out_seg,
    unsigned long long stream_ptr);

extern "C" void launch_final_write_kernel(
    unsigned long long in_vals_ptr,
    unsigned long long in_idxs_ptr,
    unsigned long long out_vals_ptr,
    unsigned long long out_idxs_ptr,
    int batch, int k,
    unsigned long long stream_ptr);
"""

# ---------------------------------------------------------------------------
# JIT-compile
# ---------------------------------------------------------------------------
_mod = None


def _module():
    global _mod
    if _mod is None:
        _mod = load_inline(
            name="topk_bitonic",
            cpp_sources=[CPP_SRC],
            cuda_sources=[CUDA_SRC],
            functions=[
                "launch_segment_kernel",
                "launch_merge_kernel",
                "launch_final_write_kernel",
            ],
            extra_cuda_cflags=[
                "-O3",
                "--use_fast_math",
                "-lineinfo",
            ],
            extra_include_paths=[
                "/usr/local/cuda-13.1/targets/x86_64-linux/include/cccl",
            ],
            verbose=False,
        )
    return _mod


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


def _merge_tree(mod, s, batch, k, seg_vals, seg_idxs, num_seg):
    """Iteratively merge segments until num_seg == 1 per row."""
    cur_vals = seg_vals
    cur_idxs = seg_idxs
    cur_num_seg = num_seg

    while cur_num_seg > 1:
        group_size = min(cur_num_seg, 1024 // k)
        if group_size < 2:
            group_size = 2

        new_num_seg = (cur_num_seg + group_size - 1) // group_size

        new_vals = torch.empty(batch, new_num_seg, k, dtype=torch.float32, device=cur_vals.device)
        new_idxs = torch.empty(batch, new_num_seg, k, dtype=torch.int64, device=cur_idxs.device)

        mod.launch_merge_kernel(
            cur_vals.data_ptr(),
            cur_idxs.data_ptr(),
            new_vals.data_ptr(),
            new_idxs.data_ptr(),
            batch, k, cur_num_seg, group_size, new_num_seg,
            s,
        )

        cur_vals = new_vals
        cur_idxs = new_idxs
        cur_num_seg = new_num_seg

    return cur_vals, cur_idxs


# ---------------------------------------------------------------------------
# Model
# ---------------------------------------------------------------------------
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))

    def forward(self, x: torch.Tensor):
        batch, n, k = self.batch, self.n, self.k
        assert x.shape == (batch, n)
        assert x.dtype == torch.float32
        assert x.is_cuda

        mod = _module()
        s = torch.cuda.current_stream().cuda_stream

        # Phase 1: segment radix sort (CUB)
        SEG_SIZE = 1024
        num_seg = (n + SEG_SIZE - 1) // SEG_SIZE

        seg_vals = torch.empty(batch, num_seg, k, dtype=torch.float32, device=x.device)
        seg_idxs = torch.empty(batch, num_seg, k, dtype=torch.int64, device=x.device)

        mod.launch_segment_kernel(
            x.data_ptr(),
            seg_vals.data_ptr(),
            seg_idxs.data_ptr(),
            batch, n, k, num_seg,
            s,
        )

        # Phase 2: merge tree
        merged_vals, merged_idxs = _merge_tree(mod, s, batch, k, seg_vals, seg_idxs, num_seg)
        # merged_vals shape: (batch, 1, k), sorted ascending

        # Phase 3: write descending output
        out_vals = torch.empty(batch, k, dtype=torch.float32, device=x.device)
        out_idxs = torch.empty(batch, k, dtype=torch.int64, device=x.device)

        mod.launch_final_write_kernel(
            merged_vals.data_ptr(),
            merged_idxs.data_ptr(),
            out_vals.data_ptr(),
            out_idxs.data_ptr(),
            batch, k,
            s,
        )

        return out_vals, out_idxs


# ---------------------------------------------------------------------------
# Module-level shims
# ---------------------------------------------------------------------------
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]