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

_CUDA_SRC = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <c10/cuda/CUDAStream.h>
#include <math_constants.h>

#define KMAX 64

__device__ __forceinline__ unsigned int fkey(float f) {
    unsigned int b = __float_as_uint(f);
    return b ^ ((b >> 31) ? 0xFFFFFFFFu : 0x80000000u);
}

// Block-level top-k of (sv,si)[0..len). Writes k sorted-descending results to out_val/out_idx.
__device__ void block_topk(float* sv, int* si, int len, int k,
                           float* out_val, long* out_idx) {
    __shared__ unsigned int hist[256];
    __shared__ int sc[256];
    __shared__ unsigned int s_prefix;
    __shared__ int s_need, s_base, s_cgt, s_ceq;
    __shared__ float topv[KMAX];
    __shared__ int   topi[KMAX];

    int T = blockDim.x, tid = threadIdx.x;
    if (tid == 0) { s_prefix = 0; s_need = k; }
    __syncthreads();

    for (int pass = 0; pass < 4; pass++) {
        int shift = 24 - 8 * pass;
        unsigned int curmask = (shift + 8 >= 32) ? 0u : (0xFFFFFFFFu << (shift + 8));
        for (int b = tid; b < 256; b += T) hist[b] = 0;
        __syncthreads();
        unsigned int prefix = s_prefix;
        for (int i = tid; i < len; i += T) {
            unsigned int u = fkey(sv[i]);
            if ((u & curmask) == prefix) {
                unsigned int d = (u >> shift) & 0xFFu;
                unsigned int active = __activemask();
                unsigned int same = __match_any_sync(active, d);
                int leader = __ffs(same) - 1;
                int cnt = __popc(same);
                if ((threadIdx.x & 31) == leader) atomicAdd(&hist[d], (unsigned)cnt);
            }
        }
        // suffix-sum of hist into sc[d] = sum_{j>=d} hist[j], in parallel
        if (tid < 256) sc[tid] = (int)hist[tid];
        __syncthreads();
        for (int off = 1; off < 256; off <<= 1) {
            int add = 0;
            if (tid < 256 && tid + off < 256) add = sc[tid + off];
            __syncthreads();
            if (tid < 256) sc[tid] += add;
            __syncthreads();
        }
        int need = s_need;
        if (tid < 256 && sc[tid] >= need && (tid == 255 || sc[tid + 1] < need)) {
            s_prefix = prefix | ((unsigned int)tid << shift);
            s_need = need - (tid == 255 ? 0 : sc[tid + 1]);
        }
        __syncthreads();
    }

    if (tid == 0) { s_base = k - s_need; s_cgt = 0; s_ceq = 0; }
    __syncthreads();
    unsigned int u_thr = s_prefix;
    int base = s_base, need = s_need;
    for (int i = tid; i < len; i += T) {
        unsigned int u = fkey(sv[i]);
        if (u > u_thr) {
            int p = atomicAdd(&s_cgt, 1);
            if (p < base) { topv[p] = sv[i]; topi[p] = si[i]; }
        } else if (u == u_thr) {
            int p = atomicAdd(&s_ceq, 1);
            if (p < need) { topv[base + p] = sv[i]; topi[base + p] = si[i]; }
        }
    }
    __syncthreads();

    int KP = 1; while (KP < k) KP <<= 1;
    for (int i = tid + k; i < KP; i += T) { topv[i] = -CUDART_INF_F; topi[i] = 0; }
    __syncthreads();
    for (int kk = 2; kk <= KP; kk <<= 1)
        for (int j = kk >> 1; j > 0; j >>= 1) {
            for (int i = tid; i < KP; i += T) {
                int ixj = i ^ j;
                if (ixj > i) {
                    bool up = ((i & kk) == 0);
                    float a = topv[i], b = topv[ixj];
                    if ((a > b) == up) {
                        topv[i] = b; topv[ixj] = a;
                        int t = topi[i]; topi[i] = topi[ixj]; topi[ixj] = t;
                    }
                }
            }
            __syncthreads();
        }
    for (int j = tid; j < k; j += T) {
        out_val[j] = topv[KP - 1 - j];
        out_idx[j] = topi[KP - 1 - j];
    }
}

__global__ void selectk_kernel(const float* __restrict__ x, float* __restrict__ ov,
                               long* __restrict__ oi, int n, int k, int chunk, int blocks_per_row) {
    extern __shared__ char smem[];
    float* sv = (float*)smem;
    int* si = (int*)(sv + chunk);
    int row = blockIdx.x / blocks_per_row;
    int cid = blockIdx.x % blocks_per_row;
    int cstart = cid * chunk;
    int len = chunk;
    if (cstart + len > n) len = n - cstart;
    const float* xr = x + (long)row * n + cstart;
    for (int i = threadIdx.x; i < len; i += blockDim.x) { sv[i] = xr[i]; si[i] = cstart + i; }
    __syncthreads();
    long outbase = ((long)row * blocks_per_row + cid) * k;
    block_topk(sv, si, len, k, ov + outbase, oi + outbase);
}

__global__ void merge_kernel(const float* __restrict__ cv, const long* __restrict__ ci,
                             float* __restrict__ ov, long* __restrict__ oi, int m, int k) {
    extern __shared__ char smem[];
    float* sv = (float*)smem;
    int* si = (int*)(sv + m);
    int row = blockIdx.x;
    const float* c = cv + (long)row * m;
    const long* d = ci + (long)row * m;
    for (int i = threadIdx.x; i < m; i += blockDim.x) { sv[i] = c[i]; si[i] = (int)d[i]; }
    __syncthreads();
    block_topk(sv, si, m, k, ov + (long)row * k, oi + (long)row * k);
}

static int next_pow2(int x) { int p = 1; while (p < x) p <<= 1; return p; }

static void set_smem(const void* fn, int bytes) {
    if (bytes > 48 * 1024)
        cudaFuncSetAttribute((void*)fn, cudaFuncAttributeMaxDynamicSharedMemorySize, bytes);
}

std::vector<torch::Tensor> topk_forward(torch::Tensor x, int64_t k) {
    int batch = x.size(0);
    int n = x.size(1);
    auto stream = at::cuda::getCurrentCUDAStream();
    auto opt_v = torch::TensorOptions().dtype(torch::kFloat32).device(x.device());
    auto opt_i = torch::TensorOptions().dtype(torch::kInt64).device(x.device());
    torch::Tensor out_val = torch::empty({batch, (long)k}, opt_v);
    torch::Tensor out_idx = torch::empty({batch, (long)k}, opt_i);

    const int SHARED_CAP_ELEMS = 24 * 1024;
    int chunk, blocks_per_row;
    if (n <= SHARED_CAP_ELEMS) { chunk = n; blocks_per_row = 1; }
    else { chunk = 2048; blocks_per_row = (n + chunk - 1) / chunk; }

    int T = 256;
    int smem = chunk * (sizeof(float) + sizeof(int));

    if (blocks_per_row == 1) {
        set_smem((const void*)selectk_kernel, smem);
        selectk_kernel<<<batch, T, smem, stream>>>(
            x.data_ptr<float>(), out_val.data_ptr<float>(), out_idx.data_ptr<long>(),
            n, k, chunk, 1);
    } else {
        torch::Tensor cv = torch::empty({batch, (long)blocks_per_row * (long)k}, opt_v);
        torch::Tensor ci = torch::empty({batch, (long)blocks_per_row * (long)k}, opt_i);
        set_smem((const void*)selectk_kernel, smem);
        selectk_kernel<<<batch * blocks_per_row, T, smem, stream>>>(
            x.data_ptr<float>(), cv.data_ptr<float>(), ci.data_ptr<long>(),
            n, k, chunk, blocks_per_row);
        int m = blocks_per_row * k;
        int Tm = 256;
        int smem_m = m * (sizeof(float) + sizeof(int));
        set_smem((const void*)merge_kernel, smem_m);
        merge_kernel<<<batch, Tm, smem_m, stream>>>(
            cv.data_ptr<float>(), ci.data_ptr<long>(),
            out_val.data_ptr<float>(), out_idx.data_ptr<long>(), m, k);
    }
    return {out_val, out_idx};
}
"""

_CPP_SRC = "std::vector<torch::Tensor> topk_forward(torch::Tensor x, int64_t k);"

_mod = load_inline(
    name="topk_radix_v2",
    cpp_sources=_CPP_SRC,
    cuda_sources=_CUDA_SRC,
    functions=["topk_forward"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=False,
)


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):
        return _mod.topk_forward(x.contiguous(), self.k)


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


def get_init_inputs():
    return [64, 8192, 8]