"""Top-k on a NeuronCore through nkilib's ``rotational_topk`` (issue #27).

``torch.topk`` with a large ``k`` does not compile on the device (neuronx-cc NCC_IXCG857 for k=4096 over the
129280-token vocab), so the eager sampler fell back to slow paths. nkilib's rotational top-k does k=4096 over
the vocab in ~8.4 ms with values identical to ``torch.topk`` (equal values may come out in another order).
"""

from __future__ import annotations

import functools

import torch

#: nkilib's tested range; larger widths go through the torch path.
_MAX_VOCAB = 151936
_MIN_K = 1024  # the sampler's force_topk; the indexer's k=512 compiles with torch.topk


def supported(scores: torch.Tensor, k: int) -> bool:
    return (
        scores.device.type == "neuron"
        and scores.dim() == 2
        and scores.dtype == torch.float32
        and _MIN_K <= k <= 4096
        and k < scores.shape[1] <= _MAX_VOCAB
    )


@functools.cache
def _kernel_and_config(rows: int, vocab: int, k: int):  # noqa: ANN202
    import nki.language as nl
    from nkilib.core.topk.rotational_topk import rotational_topk
    from nkilib.core.topk.rotational_topk_utils import create_rotational_topk_config, create_topk_config
    from torch_neuronx.nki_hop import wrap_nki

    tc = create_topk_config(inp_shape=(rows, vocab), inp_dtype=nl.float32, k=k, sorted=True, num_programs=1 if rows == 1 else 2)
    return wrap_nki(rotational_topk), create_rotational_topk_config(inp_shape=(rows, vocab), topk_config=tc)


def topk(scores: torch.Tensor, k: int) -> torch.Tensor:
    """Indices ``[rows, k]`` int64 of the ``k`` largest values of each row, in descending value order."""
    rows, vocab = scores.shape
    kernel, config = _kernel_and_config(rows, vocab, k)
    _, idx = kernel(scores.contiguous(), config)
    return idx.reshape(rows, -1)[:, :k].to(torch.int64)