"""Mirrors ``tilelang_ops/asc/sampler.py`` in pure torch.

Implemented (used by plain sampling, including nextn=0 decode): prepare_verification_metadata_tl,
expand_sampling_params_tl, map_lookup_tl.

MTP/DSpark verification only, out of scope on Neuron basic mode (NEURON_PORT.md "Scope"; the factories keep
the asc signatures and raise): remap_draft_topk_tl, get_accept_result_tl, finalize_verification_tl.
"""

from typing import Any, Literal

import torch

from ..common import DRAFT_POS_HASH_MUL
from ._unsupported import unsupported


def prepare_verification_metadata_tl(num_steps: int) -> Any:
    """kernel(ftp, seqlen_q, row_to_req, local_step, is_verify); rows are packed per request from ftp."""

    def _prepare_verification_metadata(
        ftp: torch.Tensor,
        seqlen_q: torch.Tensor,
        row_to_req: torch.Tensor,
        local_step: torch.Tensor,
        is_verify: torch.Tensor,
    ) -> None:
        total_q = row_to_req.shape[0]
        req_cnt = ftp.shape[0]
        device = ftp.device
        # row -> request without searchsorted/host sync: mark each request start, then prefix-count.
        starts = torch.zeros(total_q + 1, dtype=torch.int32, device=device)
        starts.scatter_add_(0, ftp.to(torch.int64).clamp(0, total_q), torch.ones(req_cnt, dtype=torch.int32, device=device))
        req = starts[:total_q].cumsum(0, dtype=torch.int32) - 1
        req_idx = req.clamp(min=0).to(torch.int64)
        step = torch.arange(total_q, dtype=torch.int32, device=device) - ftp.to(torch.int32)[req_idx]
        row_to_req.copy_(req)
        local_step.copy_(step)
        is_verify.copy_(step + 1 < seqlen_q.to(torch.int32)[req_idx])

    return _prepare_verification_metadata


def expand_sampling_params_tl() -> Any:
    """kernel(temperatures, top_ps, top_ks, sample_positions, seq_seeds, row_to_req, local_step,
    temperatures_q, top_ps_q, top_ks_q, sample_positions_q, seq_seeds_q)."""

    def _expand_sampling_params(
        temperatures: torch.Tensor,
        top_ps: torch.Tensor,
        top_ks: torch.Tensor,
        sample_positions: torch.Tensor,
        seq_seeds: torch.Tensor,
        row_to_req: torch.Tensor,
        local_step: torch.Tensor,
        temperatures_q: torch.Tensor,
        top_ps_q: torch.Tensor,
        top_ks_q: torch.Tensor,
        sample_positions_q: torch.Tensor,
        seq_seeds_q: torch.Tensor,
    ) -> None:
        req = row_to_req.to(torch.int64)
        temperatures_q.copy_(temperatures[req])
        top_ps_q.copy_(top_ps[req])
        top_ks_q.copy_(top_ks[req])
        # int32 arithmetic with wraparound, as in the kernel.
        sample_positions_q.copy_(sample_positions.to(torch.int32)[req] + local_step.to(torch.int32) * DRAFT_POS_HASH_MUL)
        seq_seeds_q.view(torch.int64).copy_(seq_seeds.view(torch.int64)[req])

    return _expand_sampling_params


def map_lookup_tl(
    force_topk: int,
    mode: Literal["by_key", "by_index"] = "by_key",
    value_dtype: torch.dtype = torch.float32,
) -> Any:
    """kernel(key, value, query, out) with out/query [rows, 1].

    by_index: ``value[row, query]`` or the out-of-bounds value (-inf float / -1 int).
    by_key: the value whose key equals the query (keys are unique top-k indices; with duplicate keys this
    returns their max), -inf when missing."""
    assert mode in ("by_key", "by_index")
    assert mode != "by_key" or value_dtype.is_floating_point, "by_key requires floating-point values"
    oob_value = float("-inf") if value_dtype.is_floating_point else -1

    def _map_lookup(
        key: torch.Tensor | None,
        value: torch.Tensor,
        query: torch.Tensor,
        out: torch.Tensor,
    ) -> None:
        q = query[:, 0]
        if mode == "by_index":
            q64 = q.to(torch.int64)
            in_bounds = (q64 >= 0) & (q64 < force_topk)
            got = value[:, :force_topk].gather(1, torch.where(in_bounds, q64, torch.zeros_like(q64)).unsqueeze(1)).squeeze(1)
            result = torch.where(in_bounds, got, torch.full_like(got, oob_value))
        else:
            match = key[:, :force_topk] == q.unsqueeze(1)
            v = value[:, :force_topk]
            result = torch.where(match, v, torch.full_like(v, oob_value)).amax(dim=1)
        out[:, 0].copy_(result)

    return _map_lookup


def remap_draft_topk_tl(force_topk: int, vocab_size: int) -> Any:
    unsupported("remap_draft_topk_tl")


def get_accept_result_tl(has_golden: bool) -> Any:
    unsupported("get_accept_result_tl")


def finalize_verification_tl(
    has_sum_smaller_probs: bool,
    has_confidence: bool,
    num_steps: int,
    num_stop_tokens: int,
) -> Any:
    unsupported("finalize_verification_tl")