"""Top-K kernel via bitonic top-k with in-register merge.

Approach:
  1. Pack (value, index) into int64 with a monotonic float encoding so that
     int64 descending order == float descending order.
  2. Per row: load BLOCK_N elements, find top-K, then iteratively merge with
     a running top-K buffer (all in registers via tl.cat + tl.topk). After
     all chunks, unpack and write the result.

Shapes:
  - (1, 131072, 64): single row, 8 chunks of 16384, in-register merge.
  - (64, 8192, 8): one CTA per row, one chunk of 8192.
  - (32, 16384, 32): one CTA per row, one chunk of 16384.
  - (16, 12000, 16): one CTA per row, one chunk of 16384 (with mask).
  - (128, 4096, 1): one CTA per row, one chunk of 4096.
"""
import math

import torch
import torch.nn as nn
import triton
import triton.language as tl


@triton.jit
def _pack(v, i):
    """Pack (float value, int index) into int64.

    Float-to-uint32 monotonic encoding:
      - For positive floats, set the sign bit (bits | 0x80000000).
      - For negative floats, invert all bits (~bits).
    This makes larger floats have larger unsigned keys.

    Then uint32 -> int32 by XOR with 0x80000000 to make signed comparison
    work; pack int32 in the high 32 bits of int64 and the index in the low
    32 bits. int64 signed descending order == float descending order.
    """
    bits = v.to(tl.uint32, bitcast=True)
    sign = bits & 0x80000000
    key = tl.where(sign != 0, ~bits, bits | 0x80000000)
    m = (key ^ 0x80000000).to(tl.int32)
    m_ext = m.to(tl.int64)
    return (m_ext << 32) | i.to(tl.int64)


@triton.jit
def _unpack(packed):
    """Unpack int64 into (float value, int64 index)."""
    idx = (packed & 0xFFFFFFFF).to(tl.int64)
    m_back = (packed >> 32).to(tl.int64)  # avoid int32 overflow on XOR
    key = (m_back ^ 0x80000000).to(tl.uint32)
    sign_back = key & 0x80000000
    # If sign_back set: original was positive, bits = key ^ 0x80000000
    # If sign_back not set: original was negative, bits = ~key
    bits = tl.where(sign_back != 0, key ^ 0x80000000, ~key)
    val = bits.to(tl.float32, bitcast=True)
    return val, idx


@triton.jit
def topk_kernel(
    x_ptr, val_ptr, idx_ptr,
    N,
    K: tl.constexpr,
    BLOCK_N: tl.constexpr,
    N_CHUNKS,
):
    """One CTA per row. Iteratively loads chunks, finds local top-K, and
    merges with a running top-K buffer kept in registers."""
    pid = tl.program_id(0)
    base = pid * N

    NEG_INF_F = float('-inf')
    NEG_INF_I64: tl.constexpr = -(1 << 63)  # int64 min

    running_packed = tl.full((K,), NEG_INF_I64, tl.int64)

    for chunk_idx in range(N_CHUNKS):
        chunk_start = chunk_idx * BLOCK_N
        offs = chunk_start + tl.arange(0, BLOCK_N)
        mask = offs < N
        v = tl.load(x_ptr + base + offs, mask=mask, other=NEG_INF_F)
        i = offs.to(tl.int64)

        packed = _pack(v, i)
        chunk_top = tl.topk(packed, K)

        # Merge: cat [running, chunk_top] (size 2K) -> topk K
        merged = tl.cat(running_packed, chunk_top, can_reorder=True)
        running_packed = tl.topk(merged, K)

    val, idx = _unpack(running_packed)
    k_offs = tl.arange(0, K)
    tl.store(val_ptr + pid * K + k_offs, val)
    tl.store(idx_ptr + pid * K + k_offs, idx)


def next_pow2(x):
    if x <= 1:
        return 1
    return 1 << (x - 1).bit_length()


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


class Model(nn.Module):
    """Top-k over the last dim of a 2D tensor.

    Input:
        x: (batch, n) fp32
    Output:
        values:  (batch, k) fp32, sorted descending
        indices: (batch, k) int64, into the last dim of x
    """

    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):
        B, N = x.shape
        K = self.k

        MAX_BLOCK_N = 16384
        BLOCK_N = min(MAX_BLOCK_N, next_pow2(N))
        N_CHUNKS = (N + BLOCK_N - 1) // BLOCK_N

        val = torch.empty(B, K, dtype=torch.float32, device=x.device)
        idx = torch.empty(B, K, dtype=torch.int64, device=x.device)

        if BLOCK_N <= 2048:
            num_warps = 4
        elif BLOCK_N <= 8192:
            num_warps = 8
        else:
            num_warps = 16

        grid = (B,)
        topk_kernel[grid](
            x, val, idx,
            N, K, BLOCK_N, N_CHUNKS,
            num_warps=num_warps,
        )

        return val, idx


# Module-level shims rebuilt by check.py / benchmark.py per shape.
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]