"""The pass that attaches GPU sampler functions to the IRModule."""

from typing import Dict  # noqa: UP035

import tvm
from tvm import IRModule, relax, te, tirx
from tvm.relax.frontend import nn
from tvm.script import tirx as T

from mlc_llm.op.batch_spec_verify import batch_spec_verify
from mlc_llm.op.top_p_pivot import top_p_pivot, top_p_renorm


@tvm.transform.module_pass(opt_level=0, name="AttachGPUSamplingFunc")
class AttachGPUSamplingFunc:
    """Attach GPU sampling functions to IRModule."""

    def __init__(self, target: tvm.target.Target, variable_bounds: Dict[str, int]):  # noqa: UP006
        # Specifically for RWKV workloads, which contains -1 max_seq_len
        max_batch_size = variable_bounds["batch_size"]
        self.variable_bounds = {
            "batch_size": max_batch_size,
            "num_samples": max_batch_size,
            "num_positions": 6 * max_batch_size,
        }
        self.non_negative_var = ["vocab_size"]
        self.target = target

    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
        """Entrypoint"""
        target_kind = self.target.kind.name
        if target_kind not in ["cuda", "vulkan", "metal", "webgpu"]:
            # Only enable GPU sampling for CUDA, Vulkan, Metal, and WebGPU.
            return mod

        bb = relax.BlockBuilder(mod)
        if target_kind == "webgpu":
            # Only attach functions that do not contain i8s for WebGPU
            gv_names = [
                gv.name_hint
                for gv in [
                    _attach_argsort_func(bb),
                    _attach_sample_with_top_p(bb),
                ]
            ]
        else:
            gv_names = [
                gv.name_hint
                for gv in [
                    _attach_multinomial_sampling_func(bb),
                    _attach_argsort_func(bb),
                    _attach_sample_with_top_p(bb),
                    _attach_take_probs_func(bb),
                    _attach_batch_verifier(bb),
                    _attach_renormalize_by_top_p(bb, self.target),
                ]
            ]

        mod = bb.finalize()
        for gv_name in gv_names:
            mod[gv_name] = (
                mod[gv_name]
                .with_attr("tir_var_upper_bound", self.variable_bounds)
                .with_attr("tir_non_negative_var", self.non_negative_var)
            )
        return mod


def _attach_multinomial_sampling_func(bb: relax.BlockBuilder):
    batch_size = tirx.Var("batch_size", "int64")
    num_samples = tirx.Var("num_samples", "int64")
    vocab_size = tirx.Var("vocab_size", "int64")
    probs = relax.Var("probs", relax.TensorType((batch_size, vocab_size), "float32"))
    uniform_samples = relax.Var("uniform_samples", relax.TensorType((num_samples,), "float32"))
    sample_indices = relax.Var("sample_indices", relax.TensorType((num_samples,), "int32"))
    with bb.function("multinomial_from_uniform", [probs, uniform_samples, sample_indices]):
        with bb.dataflow():
            sample_shape = relax.ShapeExpr([num_samples, 1])
            probs_tensor = nn.wrap_nested(probs, name="probs")
            uniform_samples_tensor = nn.wrap_nested(
                relax.call_pure_packed(
                    "vm.builtin.reshape",
                    uniform_samples,
                    sample_shape,
                    ty_args=relax.TensorType(sample_shape, "float32"),
                ),
                name="uniform_samples",
            )
            sample_indices_tensor = nn.wrap_nested(
                relax.call_pure_packed(
                    "vm.builtin.reshape",
                    sample_indices,
                    sample_shape,
                    ty_args=relax.TensorType(sample_shape, "int32"),
                ),
                name="sample_indices",
            )
            result_tensor = nn.multinomial_from_uniform(
                probs_tensor,
                uniform_samples_tensor,
                sample_indices_tensor,
                "int32",
                name="nn_multinomial_from_uniform",
            )
            result = bb.emit(
                relax.call_pure_packed(
                    "vm.builtin.reshape",
                    result_tensor._expr,
                    sample_indices.ty.shape,
                    ty_args=sample_indices.ty,
                )
            )
            output = bb.emit_output(result)
        gv = bb.emit_func_output(output)
    return gv


def _attach_argsort_func(bb: relax.BlockBuilder):
    batch_size = tirx.Var("batch_size", "int64")
    vocab_size = tirx.Var("vocab_size", "int64")
    probs = relax.Var("probs", relax.TensorType((batch_size, vocab_size), "float32"))
    with bb.function("argsort_probs", [probs]):
        with bb.dataflow():
            sorted_indices = bb.emit(relax.op.argsort(probs, descending=True, dtype="int32"))
            sorted_values = bb.emit_te(
                lambda unsorted_probs, sorted_indices: te.compute(
                    (batch_size, vocab_size),
                    lambda i, j: unsorted_probs[i, sorted_indices[i, j]],
                    name="take_sorted_probs",
                ),
                probs,
                sorted_indices,
                primfunc_name_hint="take_sorted_probs",
            )
            output = bb.emit_output((sorted_values, sorted_indices))
        gv = bb.emit_func_output(output)
    return gv


@T.prim_func(s_tir=True)
def full(var_result: T.handle, value: T.int32):
    """The filling function for top k."""
    batch_size = T.int32()
    result = T.match_buffer(var_result, (batch_size, 1), "int32")
    for i in T.serial(batch_size):
        with T.sblock("block"):
            vi = T.axis.spatial(batch_size, i)
            result[vi, 0] = value


def _attach_sample_with_top_p(bb: relax.BlockBuilder):
    batch_size = tirx.Var("batch_size", "int64")
    num_samples = tirx.Var("num_samples", "int64")
    vocab_size = tirx.Var("vocab_size", "int64")
    sorted_probs = relax.Var("sorted_probs", relax.TensorType((batch_size, vocab_size), "float32"))
    sorted_indices = relax.Var(
        "sorted_indices", relax.TensorType((batch_size, vocab_size), "int32")
    )
    uniform_samples = relax.Var("uniform_samples", relax.TensorType((num_samples,), "float32"))
    sample_indices = relax.Var("sample_indices", relax.TensorType((num_samples,), "int32"))
    top_p = relax.Var("top_p", relax.TensorType((batch_size,), "float32"))

    with bb.function(
        "sample_with_top_p",
        [sorted_probs, sorted_indices, uniform_samples, sample_indices, top_p],
    ):
        with bb.dataflow():
            sample_shape = relax.ShapeExpr([num_samples, 1])
            top_p_shape = relax.ShapeExpr([batch_size, 1])
            sorted_probs_tensor = nn.wrap_nested(sorted_probs, name="sorted_probs")
            sorted_indices_tensor = nn.wrap_nested(sorted_indices, name="sorted_indices")
            uniform_samples_tensor = nn.wrap_nested(
                relax.call_pure_packed(
                    "vm.builtin.reshape",
                    uniform_samples,
                    sample_shape,
                    ty_args=relax.TensorType(sample_shape, "float32"),
                ),
                name="uniform_samples",
            )
            sample_indices_tensor = nn.wrap_nested(
                relax.call_pure_packed(
                    "vm.builtin.reshape",
                    sample_indices,
                    sample_shape,
                    ty_args=relax.TensorType(sample_shape, "int32"),
                ),
                name="sample_indices",
            )
            top_p_tensor = nn.wrap_nested(
                relax.call_pure_packed(
                    "vm.builtin.reshape",
                    top_p,
                    top_p_shape,
                    ty_args=relax.TensorType(top_p_shape, "float32"),
                ),
                name="sample_indices",
            )
            top_k_tensor = nn.tensor_ir_op(
                full,
                name_hint="full",
                args=[vocab_size],
                out=nn.Tensor.placeholder(
                    [batch_size, 1],
                    "int32",
                ),
            )

            result_tensor = nn.sample_top_p_top_k_from_sorted_prob(
                sorted_probs_tensor,
                sorted_indices_tensor,
                top_p_tensor,
                top_k_tensor,
                uniform_samples_tensor,
                sample_indices_tensor,
            )
            result = bb.emit_output(
                relax.call_pure_packed(
                    "vm.builtin.reshape",
                    result_tensor._expr,
                    sample_indices.ty.shape,
                    ty_args=sample_indices.ty,
                )
            )
        gv = bb.emit_func_output(result)
    return gv


def _attach_renormalize_by_top_p(bb: relax.BlockBuilder, target: tvm.target.Target):
    batch_size = tirx.Var("batch_size", "int64")
    vocab_size = tirx.Var("vocab_size", "int64")
    num_pivots = 3
    probs = relax.Var("probs", relax.TensorType((batch_size, vocab_size), "float32"))
    top_p = relax.Var("top_p", relax.TensorType((batch_size,), "float32"))
    init_pivots = relax.Var("init_pivots", relax.TensorType((batch_size, num_pivots), "float32"))
    with bb.function("renormalize_by_top_p", [probs, top_p, init_pivots]):
        with bb.dataflow():
            cutoff_output = bb.emit(
                relax.call_tir(
                    bb.add_func(top_p_pivot(num_pivots, target), "top_p_pivot_cutoff"),
                    args=[probs, top_p, init_pivots],
                    out_ty=[top_p.ty, top_p.ty],
                )
            )
            final_pivot = cutoff_output[0]
            renorm_sum = cutoff_output[1]
            renormalized_probs = bb.emit_output(
                relax.call_tir(
                    bb.add_func(top_p_renorm(target), "top_p_renorm_after_cutoff"),
                    args=[probs, final_pivot, renorm_sum],
                    out_ty=probs.ty,
                )
            )
        gv = bb.emit_func_output(renormalized_probs)
    return gv


def _attach_take_probs_func(bb: relax.BlockBuilder):
    batch_size = tirx.Var("batch_size", "int64")
    num_samples = tirx.Var("num_samples", "int64")
    num_positions = tirx.Var("num_positions", "int64")
    vocab_size = tirx.Var("vocab_size", "int64")
    unsorted_probs = relax.Var(
        "unsorted_probs", relax.TensorType((batch_size, vocab_size), "float32")
    )
    sorted_indices = relax.Var(
        "sorted_indices", relax.TensorType((batch_size, vocab_size), "int32")
    )
    sample_indices = relax.Var("sample_indices", relax.TensorType((num_samples,), "int32"))
    sampling_results = relax.Var("sampling_result", relax.TensorType((num_samples,), "int32"))
    top_prob_offsets = relax.Var("lobprob_offsets", relax.TensorType((num_positions,), "int32"))

    @T.prim_func(s_tir=True)
    def sampler_take_probs_tir(
        var_unsorted_probs: T.handle,
        var_sorted_indices: T.handle,
        var_sample_indices: T.handle,
        var_sampling_results: T.handle,
        var_top_prob_offsets: T.handle,
        var_sampled_values: T.handle,
        var_top_prob_probs: T.handle,
        var_top_prob_indices: T.handle,
    ):
        batch_size = T.int32()
        num_samples = T.int32()
        num_positions = T.int32()
        vocab_size = T.int32()
        unsorted_probs = T.match_buffer(var_unsorted_probs, (batch_size, vocab_size), "float32")
        sorted_indices = T.match_buffer(var_sorted_indices, (batch_size, vocab_size), "int32")
        sample_indices = T.match_buffer(var_sample_indices, (num_samples,), "int32")
        sampling_results = T.match_buffer(var_sampling_results, (num_samples,), "int32")
        top_prob_offsets = T.match_buffer(var_top_prob_offsets, (num_positions,), "int32")
        sampled_values = T.match_buffer(var_sampled_values, (num_samples,), "float32")
        top_prob_probs = T.match_buffer(var_top_prob_probs, (num_positions,), "float32")
        top_prob_indices = T.match_buffer(var_top_prob_indices, (num_positions,), "int32")
        for i in T.serial(num_positions):
            with T.sblock("top_prob"):
                vi = T.axis.spatial(num_positions, i)
                # Reads are data-dependent gathers; declare full-buffer read
                # regions explicitly so tirx does not infer data-dependent regions.
                T.reads(
                    top_prob_offsets[vi],
                    sorted_indices[0:batch_size, 0:vocab_size],
                    unsorted_probs[0:batch_size, 0:vocab_size],
                )
                T.writes(top_prob_indices[vi], top_prob_probs[vi])
                row = T.floordiv(top_prob_offsets[vi], vocab_size)
                col = T.floormod(top_prob_offsets[vi], vocab_size)
                top_prob_indices[vi] = sorted_indices[row, col]
                top_prob_probs[vi] = unsorted_probs[row, sorted_indices[row, col]]
        for i in T.serial(num_samples):
            with T.sblock("sample"):
                vj = T.axis.spatial(num_samples, i)
                T.reads(
                    sample_indices[vj],
                    sampling_results[vj],
                    unsorted_probs[0:batch_size, 0:vocab_size],
                )
                T.writes(sampled_values[vj])
                sampled_values[vj] = unsorted_probs[sample_indices[vj], sampling_results[vj]]

    args = [
        unsorted_probs,
        sorted_indices,
        sample_indices,
        sampling_results,
        top_prob_offsets,
    ]
    with bb.function("sampler_take_probs", args):
        with bb.dataflow():
            taken_probs_indices = bb.emit_output(
                relax.call_tir(
                    bb.add_func(sampler_take_probs_tir, "sampler_take_probs_tir"),
                    args,
                    out_ty=[
                        relax.TensorType((num_samples,), "float32"),
                        relax.TensorType((num_positions,), "float32"),
                        relax.TensorType((num_positions,), "int32"),
                    ],
                )
            )
        gv = bb.emit_func_output(taken_probs_indices)
    return gv


def _attach_batch_verifier(bb: relax.BlockBuilder):
    num_nodes = tirx.Var("num_nodes", "int64")
    nbatch = tirx.Var("nbatch", "int64")
    vocab_size = tirx.Var("vocab_size", "int64")
    draft_probs = relax.Var("draft_probs", relax.TensorType((num_nodes, vocab_size), "float32"))
    draft_tokens = relax.Var("draft_tokens", relax.TensorType((num_nodes,), "int32"))
    model_probs = relax.Var("model_probs", relax.TensorType((num_nodes, vocab_size), "float32"))
    token_tree_first_child = relax.Var(
        "token_tree_first_child", relax.TensorType((num_nodes,), "int32")
    )
    token_tree_next_sibling = relax.Var(
        "token_tree_next_sibling", relax.TensorType((num_nodes,), "int32")
    )
    uniform_samples = relax.Var("uniform_samples", relax.TensorType((num_nodes,), "float32"))
    token_tree_parent_ptr = relax.Var("token_tree_parent_ptr", relax.TensorType((nbatch,), "int32"))
    args = [
        draft_probs,
        draft_tokens,
        model_probs,
        token_tree_first_child,
        token_tree_next_sibling,
        uniform_samples,
        token_tree_parent_ptr,
    ]
    with bb.function("sampler_verify_draft_tokens", args):
        with bb.dataflow():
            res = bb.emit_output(
                relax.call_tir_inplace(
                    bb.add_func(
                        batch_spec_verify(vocab_size),
                        "batch_verify_on_gpu_single_kernel",
                    ),
                    args,
                    inplace_indices=[
                        args.index(model_probs),
                        args.index(token_tree_parent_ptr),
                    ],
                    out_ty=[
                        model_probs.ty,
                        token_tree_parent_ptr.ty,
                    ],
                )
            )
        gv = bb.emit_func_output(res)
    return gv