"""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:
    """The frontend ``get_accept_result_ref`` (static shapes: a cumsum prefix), with the kernel's argument list."""
    from ..sampler import get_accept_result_ref

    def _accept(raw_accept, golden, draft_pool, read_rows, candidates, is_verify, ftp, row_to_req, token_ids, accept, is_final):  # noqa: ANN001
        draft = draft_pool.index_select(0, read_rows.long())
        tokens, eff, final = get_accept_result_ref(
            raw_accept, golden if has_golden else None, draft, candidates, is_verify, ftp, row_to_req
        )
        token_ids.copy_(tokens)
        accept.copy_(eff)
        is_final.copy_(final)

    return _accept


def finalize_verification_tl(
    has_sum_smaller_probs: bool,
    has_confidence: bool,
    num_steps: int,
    num_stop_tokens: int,
) -> Any:
    """The frontend ``finalize_verification_ref`` with the kernel's argument list, its boolean-mask final-row lookup
    replaced by a static one (exactly one final row per request: ``ftp + last_valid_offset``)."""
    from ..sampler import compute_confidence_error_ref

    def _finalize(
        effective_accept, is_final, effective_token_ids, token_logprobs, ssp, top_lp, top_ids, ftp, seqlen_q,
        row_to_req, is_verify, accept_probs, confidence_pool, read_rows, stop_token_ids,
        last_valid_offset, num_generated_tokens, num_verified_tokens, batch_confidence_error, emitted_stop_token,
        sampled_token_ids, sampled_token_logprobs, sampled_ssp, sampled_top_lp, sampled_top_ids,
        verified_token_ids, verified_token_logprobs, verified_ssp, verified_top_lp, verified_top_ids,
    ):  # noqa: ANN001
        q = effective_accept.shape[0]
        req = row_to_req.long()
        rows = torch.arange(q, dtype=torch.int64, device=ftp.device)
        lvo = torch.zeros(ftp.shape[0], dtype=torch.int64, device=ftp.device)
        lvo.scatter_add_(0, req, effective_accept.to(torch.int64))
        last_valid_offset.copy_(lvo.to(last_valid_offset.dtype))
        final_index = ftp.long() + lvo
        local_step = rows - ftp.long()[req]
        step_index = local_step.clamp(max=num_steps - 1)
        gen = torch.zeros(num_steps, dtype=torch.int32, device=ftp.device)
        ver = torch.zeros(num_steps, dtype=torch.int32, device=ftp.device)
        gen.scatter_add_(0, step_index, effective_accept.to(torch.int32))
        ver.scatter_add_(0, step_index, is_verify.to(torch.int32))
        num_generated_tokens.copy_(gen)
        num_verified_tokens.copy_(ver)
        if has_confidence:
            rates = confidence_pool.index_select(0, read_rows.long())
            batch_confidence_error.copy_(compute_confidence_error_ref(accept_probs, rates, is_verify).reshape(1))
            row_stops = stop_token_ids.index_select(0, req)
            hits = ((effective_token_ids.unsqueeze(1) == row_stops).any(dim=1)) & (effective_accept | is_final)
            counts = torch.zeros(ftp.shape[0], dtype=torch.int32, device=ftp.device)
            counts.scatter_add_(0, req, hits.to(torch.int32))
            emitted_stop_token.copy_(counts > 0)
        else:
            batch_confidence_error.zero_()
        sampled_token_ids.copy_(effective_token_ids.index_select(0, final_index))
        sampled_token_logprobs.copy_(token_logprobs.index_select(0, final_index))
        if has_sum_smaller_probs:
            sampled_ssp.copy_(ssp.index_select(0, final_index))
            verified_ssp.copy_(torch.where(effective_accept, ssp, 0.0))
        if top_lp.shape[-1]:
            sampled_top_lp.copy_(top_lp.index_select(0, final_index))
            sampled_top_ids.copy_(top_ids.index_select(0, final_index))
            verified_top_lp.copy_(torch.where(effective_accept.unsqueeze(-1), top_lp, 0.0))
            verified_top_ids.copy_(torch.where(effective_accept.unsqueeze(-1), top_ids, 0).to(torch.int32))
        verified_token_ids.copy_(torch.where(effective_accept, effective_token_ids, -1).to(torch.int32))
        verified_token_logprobs.copy_(torch.where(effective_accept, token_logprobs, 0.0))

    return _finalize