"""Top-k over the last dim via a two-kernel selection network.

Strategy (memory-bound, a single coalesced read of the input per row):
  Tile kernel (one block per (row, tile), TILE=1024, 1024 threads):
    1. Each thread loads exactly one element of its tile into shared memory.
    2. A bitonic sort of the (<=1024) tile elements in shared memory.
    3. The top-K of the tile are written to a candidate buffer (fp32 value +
       int32 column index). Because every tile keeps its own top-K, the union
       of all tiles' top-K contains the true global top-K of the row.
  Merge kernel (one block per row, 1024 threads):
    4. Loads the row's ntiles*K candidates, bitonic-sorts them in shared
       memory, and writes out the global top-K (values + int64 indices).

The only required HBM traffic is the one input read; candidate buffers are
small (<= 64 KiB / row). No PyTorch selection / sort primitive is used.

TILE is fixed at 1024 and ntiles = ceil(n/1024), so every shape gets a healthy
number of blocks for SM occupancy and the per-tile candidates never exceed
ntiles*K <= 8192 (fits in <= 64 KiB of dynamic shared memory).
"""
import os
import sys
import sysconfig

os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "9.0")

_LOCAL_PYINC = "/workspace/problems/05_topk_bitonic/_pyinc/python3.11"


def _patch_sysconfig():
    orig = sysconfig.get_path

    def _get_path(name, *a, **k):
        p = orig(name, *a, **k)
        if p and "python3.11" in str(p) and _LOCAL_PYINC not in str(p):
            return _LOCAL_PYINC
        return p

    sysconfig.get_path = _get_path


_patch_sysconfig()

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

_TILE = 1024
_MAX_NPAD = 8192  # ntiles*K is bounded by this; shared mem = NPAD*8 bytes

_CUDA_SRC = r"""
#include <cuda_runtime.h>
#include <cstdint>

#define MAXTILE 1024

__device__ __forceinline__ float neg_inf() { return __int_as_float(0xff800000); }

// Ascending bitonic sort of arr[0..N-1] (N a power of two) in shared memory.
__device__ void bitonic(float *arr, int64_t *idx, int N) {
    for (int k = 2; k <= N; k <<= 1) {
        for (int j = k >> 1; j > 0; j >>= 1) {
            for (int i = threadIdx.x; i < N; i += blockDim.x) {
                int l = i;
                int p = l ^ j;
                if (p > l) {
                    bool dir = ((l & k) == 0);
                    float a = arr[l], b = arr[p];
                    int64_t ai = idx[l], bi = idx[p];
                    if (dir) {
                        if (a > b) { arr[l] = b; arr[p] = a; idx[l] = bi; idx[p] = ai; }
                    } else {
                        if (a < b) { arr[l] = b; arr[p] = a; idx[l] = bi; idx[p] = ai; }
                    }
                }
            }
            __syncthreads();
        }
    }
}

// Parallel k-way MERGE over the block's candidates. Each thread holds 'chunk'
// register candidates which are sorted descending once up front; each of the K
// rounds offers vv[ptr] (the thread's current head), the block reduces to the
// global max via warp shuffles + a tiny shared-memory cross-warp combine, and
// the winning thread advances its pointer. This is a correct merge of the
// per-thread sorted lists, so the emitted top-K is globally sorted descending.
__device__ void kway(int K, int chunk, float *vv, int *vi, size_t out_base,
                     float *out_val, int64_t *out_idx) {
    for (int j = 1; j < chunk; j++) {
        float key = vv[j]; int ki = vi[j];
        int p = j - 1;
        while (p >= 0 && vv[p] < key) { vv[p + 1] = vv[p]; vi[p + 1] = vi[p]; p--; }
        vv[p + 1] = key; vi[p + 1] = ki;
    }
    __shared__ float smax[32];
    __shared__ int slane[32];
    __shared__ int swin[1];
    int T = blockDim.x;
    int lane = threadIdx.x & 31;
    int wid = threadIdx.x >> 5;
    int ptr = 0;
    for (int r = 0; r < K; r++) {
        float lm = (ptr < chunk) ? vv[ptr] : neg_inf();
        int li = (ptr < chunk) ? vi[ptr] : -1;
        float mv = lm; int ml = lane;
        for (int off = 1; off < 32; off <<= 1) {
            float ov = __shfl_xor_sync(0xffffffff, mv, off);
            int ol = __shfl_xor_sync(0xffffffff, ml, off);
            if (ov > mv) { mv = ov; ml = ol; }
        }
        if (lane == 0) { smax[wid] = mv; slane[wid] = ml; }
        __syncthreads();
        int gwin = -1;
        if (threadIdx.x == 0) {
            float gm = neg_inf(); int gw = -1, gl = -1;
            int nwarps = T >> 5;
            for (int w = 0; w < nwarps; w++) {
                if (smax[w] > gm) { gm = smax[w]; gw = w; gl = slane[w]; }
            }
            gwin = gw * 32 + gl;
            swin[0] = gwin;
        }
        __syncthreads();
        gwin = swin[0];
        if (threadIdx.x == gwin) {
            out_val[out_base + r] = lm;
            out_idx[out_base + r] = (int64_t)li;
            ptr++;
        }
        __syncthreads();
    }
}

extern "C" void tile_cuda(const float *x, int batch, int n, int K, int ntiles, int TILE,
                          float *cand_val, int64_t *cand_idx);
extern "C" void merge_cuda(const float *cand_val, const int64_t *cand_idx, int batch,
                           int K, int ntiles, int M, float *out_val, int64_t *out_idx);

__global__ void tile_kernel_bitonic(const float *__restrict__ x, int batch, int n, int K, int ntiles,
                                      int TILE,
                                      float *__restrict__ cand_val, int64_t *__restrict__ cand_idx) {
    int gid = blockIdx.x;
    int row = gid / ntiles;
    int t = gid % ntiles;
    int start = t * TILE;
    int end = (start + TILE < n) ? start + TILE : n;
    __shared__ float sv[MAXTILE];
    __shared__ int64_t si[MAXTILE];
    for (int i = threadIdx.x; i < TILE; i += blockDim.x) {
        int pos = start + i;
        if (pos < end) { sv[i] = x[row * n + pos]; si[i] = pos; }
        else { sv[i] = neg_inf(); si[i] = -1; }
    }
    __syncthreads();
    bitonic(sv, si, TILE);
    __syncthreads();
    if (threadIdx.x < K) {
        int src = TILE - 1 - threadIdx.x;
        cand_val[(size_t)gid * K + threadIdx.x] = sv[src];
        cand_idx[(size_t)gid * K + threadIdx.x] = si[src];
    }
}

__global__ void tile_kernel_kway(const float *__restrict__ x, int batch, int n, int K, int ntiles,
                                  int TILE,
                                  float *__restrict__ cand_val, int64_t *__restrict__ cand_idx) {
    int gid = blockIdx.x;
    int row = gid / ntiles;
    int t = gid % ntiles;
    int start = t * TILE;
    int T = blockDim.x;
    int chunk = (TILE + T - 1) / T;
    float vv[4];
    int vi[4];
    for (int j = 0; j < chunk; j++) {
        int pos = start + threadIdx.x * chunk + j;
        if (pos < n) { vv[j] = x[row * n + pos]; vi[j] = pos; }
        else { vv[j] = neg_inf(); vi[j] = -1; }
    }
    kway(K, chunk, vv, vi, (size_t)gid * K, cand_val, cand_idx);
}

__global__ void merge_kernel(const float *__restrict__ cand_val, const int64_t *__restrict__ cand_idx,
                             int batch, int K, int ntiles, int M,
                             float *__restrict__ out_val, int64_t *__restrict__ out_idx) {
    int T = blockDim.x;
    int per = (M + T - 1) / T;
    float vv[128];
    int vi[128];
    size_t base = (size_t)blockIdx.x * M;
    for (int j = 0; j < per; j++) {
        int idx = threadIdx.x * per + j;
        if (idx < M) { vv[j] = cand_val[base + idx]; vi[j] = (int)cand_idx[base + idx]; }
        else { vv[j] = neg_inf(); vi[j] = -1; }
    }
    kway(K, per, vv, vi, (size_t)blockIdx.x * K, out_val, out_idx);
}

void tile_cuda(const float *x, int batch, int n, int K, int ntiles, int TILE,
               float *cand_val, int64_t *cand_idx) {
    if (K <= 32)
        tile_kernel_kway<<<batch * ntiles, 256>>>(x, batch, n, K, ntiles, TILE, cand_val, cand_idx);
    else
        tile_kernel_bitonic<<<batch * ntiles, 256>>>(x, batch, n, K, ntiles, TILE, cand_val, cand_idx);
}

void merge_cuda(const float *cand_val, const int64_t *cand_idx, int batch,
                int K, int ntiles, int M, float *out_val, int64_t *out_idx) {
    merge_kernel<<<batch, 128>>>(cand_val, cand_idx, batch, K, ntiles, M, out_val, out_idx);
}
"""

_CPP_SRC = r"""
extern "C" void tile_cuda(const float* x, int batch, int n, int K, int ntiles, int TILE,
                          float* cand_val, int64_t* cand_idx);
extern "C" void merge_cuda(const float* cand_val, const int64_t* cand_idx, int batch,
                           int K, int ntiles, int M, float* out_val, int64_t* out_idx);
#include <torch/extension.h>

void topk_cuda_torch(torch::Tensor x, int64_t K, torch::Tensor out_val,
                     torch::Tensor out_idx) {
    auto sizes = x.sizes();
    int batch = sizes[0];
    int n = sizes[1];
    const int TILE_TARGET = 512;   // aim for ~512-byte tiles (good SM occupancy)
    const int MMAX = 8192;         // cap ntiles*K so the merge stays cheap
    int ntiles = (n + TILE_TARGET - 1) / TILE_TARGET;
    if (ntiles * (int)K > MMAX) ntiles = MMAX / (int)K;
    // Round the actual tile size up to a power of two (bitonic needs pow2; the
    // k-way tile is also most robust with pow2 tiles).
    int need = (n + ntiles - 1) / ntiles;
    int TILE = 1;
    while (TILE < need) TILE <<= 1;
    int M = ntiles * (int)K;
    long csize = (long)batch * ntiles * K;
    // Reuse candidate buffers across calls to avoid per-call allocation.
    static torch::Tensor cand_val, cand_idx;
    if (cand_val.numel() != csize || cand_val.device() != x.device()) {
        cand_val = torch::empty({csize}, torch::dtype(torch::kF32).device(x.device()));
        cand_idx = torch::empty({csize}, torch::dtype(torch::kI64).device(x.device()));
    }
    tile_cuda(x.data_ptr<float>(), batch, n, (int)K, ntiles, TILE,
              cand_val.data_ptr<float>(), cand_idx.data_ptr<int64_t>());
    merge_cuda(cand_val.data_ptr<float>(), cand_idx.data_ptr<int64_t>(), batch, (int)K,
               ntiles, M, out_val.data_ptr<float>(), out_idx.data_ptr<int64_t>());
}
"""

_module = None


def _get_module():
    global _module
    if _module is None:
        _module = load_inline(
            name="topk_bitonic2",
            cpp_sources=_CPP_SRC,
            cuda_sources=_CUDA_SRC,
            functions=["topk_cuda_torch"],
            with_cuda=True,
            extra_cuda_cflags=["-O3", "--use_fast_math"],
        )
    return _module


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):
        mod = _get_module()
        if not hasattr(self, "_ov") or self._ov.device != x.device:
            self._ov = torch.empty(x.shape[0], self.k, dtype=torch.float32, device=x.device)
            self._oi = torch.empty(x.shape[0], self.k, dtype=torch.int64, device=x.device)
        mod.topk_cuda_torch(x, self.k, self._ov, self._oi)
        return self._ov, self._oi


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]