"""Fast input-logprob path must match the log-softmax reference.

The fast path (SGLANG_ENABLE_FAST_INPUT_LOGPROBS) computes token / top-k /
token-ids logprobs directly from logits with a per-row logsumexp normalizer,
never materializing the full-vocab log-softmax. Same math, so results must
agree with the reference path to floating-point tolerance, with identical
top-k indices, across chunk splits and heterogeneous per-sequence params.
"""

import itertools
import unittest
from types import SimpleNamespace

import torch

from sglang.srt.layers.logprob_processor import (
    InputLogprobProcessor,
    compute_row_log_normalizer,
)
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase

register_cpu_ci(est_time=30, suite="base-a-test-cpu")

VOCAB = 11
# Heterogeneous per-sequence parameters; uniform ones hide misalignment.
TOPK_CYCLE = [2, 0, 3]
# [] is a valid probe set distinct from None (opt-out).
TOKEN_IDS_CYCLE = [[0, 3], None, [1], []]


def _build_batch(seq_specs, dtype, vocab=VOCAB):
    """seq_specs: list of (extend_len, logprob_start_len). Mirrors
    LogitsProcessor._get_pruned_states for the extend-with-logprobs path."""
    pruned_rows = []
    token_to_seq_idx = []
    sample_indices = []
    input_logprob_indices = []
    pruned_lens = []
    sample_pt = -1
    lp_pt = 0
    for idx, (extend_len, start) in enumerate(seq_specs):
        eff_start = start - 1 if extend_len == start else start
        rows = extend_len - eff_start
        pruned_rows.append(torch.randn(rows, vocab).to(dtype))
        token_to_seq_idx.extend([idx] * rows)
        sample_pt += rows
        sample_indices.append(sample_pt)
        n_lp = extend_len - start
        input_logprob_indices.extend([lp_pt + i for i in range(n_lp)])
        lp_pt += rows
        pruned_lens.append(n_lp)
    metadata = SimpleNamespace(
        extend_return_top_logprob=True,
        extend_token_ids_logprob=True,
        top_logprobs_nums=[TOPK_CYCLE[i % 3] for i in range(len(seq_specs))],
        extend_logprob_pruned_lens_cpu=pruned_lens,
        extend_input_logprob_token_ids_gpu=torch.zeros(
            len(input_logprob_indices), dtype=torch.int64
        ),
        token_ids_logprobs=[
            TOKEN_IDS_CYCLE[i % len(TOKEN_IDS_CYCLE)] for i in range(len(seq_specs))
        ],
    )
    return (
        torch.cat(pruned_rows),
        torch.tensor(sample_indices, dtype=torch.int64),
        torch.tensor(input_logprob_indices, dtype=torch.int64),
        token_to_seq_idx,
        metadata,
    )


def _run(proc, batch, fast, chunk_size):
    pruned_states, sample_indices, input_logprob_indices, t2s, metadata = batch
    proc.enable_logprobs_chunk = chunk_size is not None
    proc.logprobs_chunk_size = chunk_size if chunk_size is not None else 10**9
    proc.enable_fast_input_logprobs = fast

    def get_logits_fn(states, lm_head, logits_metadata, **kwargs):
        return states

    return proc.forward(
        pruned_states=pruned_states,
        sample_indices=sample_indices,
        input_logprob_indices=input_logprob_indices,
        token_to_seq_idx=t2s,
        lm_head=None,
        get_logits_fn=get_logits_fn,
        logits_metadata=metadata,
    )


def _assert_nested_close(test, ref, got, label, rtol, atol):
    test.assertEqual(_shape_of(ref), _shape_of(got), label)
    ref_flat = _flatten(ref)
    got_flat = _flatten(got)
    if ref_flat:
        torch.testing.assert_close(
            torch.tensor(ref_flat, dtype=torch.float64),
            torch.tensor(got_flat, dtype=torch.float64),
            rtol=rtol,
            atol=atol,
            msg=label,
        )


def _flatten(nested):
    if isinstance(nested, list):
        return [x for item in nested for x in _flatten(item)]
    return [nested]


def _shape_of(nested):
    if isinstance(nested, list):
        return [_shape_of(item) for item in nested]
    return None


class TestFastInputLogprobs(CustomTestCase):
    def _sweep(self, dtype, rtol, atol):
        torch.manual_seed(0)
        proc = InputLogprobProcessor()
        # (extend_len, start); start == extend_len is the degenerate
        # zero-logprob-row shape.
        menu = [(1, 1), (3, 0), (4, 1), (5, 5), (6, 2)]
        tried = 0
        for n_seqs in (1, 2, 3):
            for combo in itertools.product(menu, repeat=n_seqs):
                batch = _build_batch(list(combo), dtype)
                for chunk_size in (None, 1, 2, 3, 5):
                    tried += 1
                    ref, ref_sampled = _run(proc, batch, False, chunk_size)
                    got, got_sampled = _run(proc, batch, True, chunk_size)
                    label = f"specs={list(combo)} chunk={chunk_size} dtype={dtype}"
                    # Top-k order comes from the same values shifted by a
                    # per-row constant, so indices must match exactly.
                    self.assertEqual(ref.top_logprobs_idx, got.top_logprobs_idx, label)
                    self.assertEqual(
                        ref.token_ids_logprobs_idx, got.token_ids_logprobs_idx, label
                    )
                    _assert_nested_close(
                        self,
                        ref.top_logprobs_val,
                        got.top_logprobs_val,
                        label,
                        rtol,
                        atol,
                    )
                    _assert_nested_close(
                        self,
                        ref.token_ids_logprobs_val,
                        got.token_ids_logprobs_val,
                        label,
                        rtol,
                        atol,
                    )
                    torch.testing.assert_close(
                        ref.token_logprobs.float(),
                        got.token_logprobs.float(),
                        rtol=rtol,
                        atol=atol,
                        msg=label,
                    )
                    torch.testing.assert_close(ref_sampled, got_sampled, msg=label)
        self.assertGreater(tried, 100)

    def test_fast_matches_reference_fp32(self):
        self._sweep(torch.float32, rtol=1e-5, atol=1e-5)

    def test_fast_matches_float64_truth_bf16(self):
        # bf16 log_softmax rounds near-ties together, so the reference path's
        # top-k ORDER is not reproducible from raw logits; validate the fast
        # path against float64 ground truth instead. The fast path only
        # rounds at the bf16 logits themselves (normalizer is fp32), so it
        # sits much closer to the truth than bf16 resolution.
        torch.manual_seed(0)
        proc = InputLogprobProcessor()
        menu = [(1, 1), (3, 0), (4, 1), (5, 5), (6, 2)]
        for n_seqs in (1, 2, 3):
            for combo in itertools.product(menu, repeat=n_seqs):
                batch = _build_batch(list(combo), torch.bfloat16)
                pruned_states, _, input_logprob_indices, _, metadata = batch
                truth = torch.log_softmax(pruned_states.double(), dim=-1)[
                    input_logprob_indices
                ]
                for chunk_size in (None, 2, 5):
                    got, _ = _run(proc, batch, True, chunk_size)
                    label = f"specs={list(combo)} chunk={chunk_size}"
                    self._assert_rows_match_truth(
                        got, truth, metadata, label, atol=1e-4
                    )

    def _assert_rows_match_truth(self, got, truth, metadata, label, atol):
        pt = 0
        for s, pruned_len in enumerate(metadata.extend_logprob_pruned_lens_cpu):
            if pruned_len <= 0:
                self.assertEqual(got.top_logprobs_val[s], [], label)
                continue
            k = metadata.top_logprobs_nums[s]
            probe_ids = metadata.token_ids_logprobs[s]
            for j in range(pruned_len):
                row_truth = truth[pt + j]
                vals = got.top_logprobs_val[s][j]
                idxs = got.top_logprobs_idx[s][j]
                self.assertEqual(len(vals), k, label)
                for v, i in zip(vals, idxs):
                    self.assertAlmostEqual(
                        v, row_truth[i].item(), delta=atol, msg=label
                    )
                if probe_ids is not None:
                    probe_vals = got.token_ids_logprobs_val[s][j]
                    for v, i in zip(probe_vals, probe_ids):
                        self.assertAlmostEqual(
                            v, row_truth[i].item(), delta=atol, msg=label
                        )
            pt += pruned_len

    def test_shift_invariant_large_offset(self):
        # Regression: a large common fp32 offset must not round the log-sum
        # term away. Uniform logits at 1e8 have true logprob -log(vocab).
        for device in ("cpu", "cuda") if torch.cuda.is_available() else ("cpu",):
            logits = torch.full((4, 1000), 1e8, dtype=torch.float32, device=device)
            row_max, row_log_sum = compute_row_log_normalizer(logits)
            logprob = (logits[:, 0].float() - row_max) - row_log_sum
            expected = -torch.log(torch.tensor(1000.0))
            torch.testing.assert_close(
                logprob.cpu(), expected.expand(4), rtol=1e-5, atol=1e-5
            )

    def test_fast_path_emits_fp32_logprobs(self):
        # Chosen dtype policy: the fast path returns fp32 token logprobs
        # regardless of logits dtype (the normalizer is fp32, so fp32 is the
        # true precision of the result), while the log_softmax path keeps
        # the logits dtype. Runs on CPU CI so the policy is pinned even
        # where the CUDA kernels never execute.
        proc = InputLogprobProcessor()
        batch = _build_batch([(4, 1), (3, 0)], torch.bfloat16)
        got, _ = _run(proc, batch, True, None)
        self.assertEqual(got.token_logprobs.dtype, torch.float32)
        ref, _ = _run(proc, batch, False, None)
        self.assertEqual(ref.token_logprobs.dtype, torch.bfloat16)

    def test_logsumexp_module_imports(self):
        # Runs on CPU CI too: catches import rot in the CUDA-only kernel
        # module, whose imports otherwise only execute on GPU machines.
        import sglang.srt.layers.logsumexp  # noqa: F401

    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
    def test_fused_tiny_shapes(self):
        from sglang.srt.layers.logsumexp import row_logsumexp_topk

        for rows, cols, k in ((1, 1, 1), (3, 2, 2), (2, 3, 1), (2, 5, 5)):
            logits = torch.randn(rows, cols, device="cuda")
            got_m, got_ls, got_v, got_i = row_logsumexp_topk(logits, k)
            ref_v, ref_i = torch.topk(logits, k, dim=-1, sorted=True)
            self.assertTrue(torch.equal(got_v, ref_v.float()), (rows, cols, k))
            self.assertTrue(torch.equal(got_i, ref_i), (rows, cols, k))
            torch.testing.assert_close(
                got_m + got_ls,
                torch.logsumexp(logits.double(), dim=-1).float(),
                rtol=1e-5,
                atol=1e-5,
            )

    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
    def test_fused_run_to_run_deterministic(self):
        from sglang.srt.layers.logsumexp import row_logsumexp_topk

        torch.manual_seed(0)
        logits = torch.randn(512, 151936, dtype=torch.bfloat16, device="cuda")
        a = row_logsumexp_topk(logits, 5)
        b = row_logsumexp_topk(logits, 5)
        self.assertTrue(all(torch.equal(p, q) for p, q in zip(a, b)))

    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
    def test_fused_logsumexp_topk_matches_torch(self):
        from sglang.srt.layers.logsumexp import row_logsumexp, row_logsumexp_topk

        torch.manual_seed(0)
        for k in (1, 2, 3, 5, 8):
            for dtype in (torch.float32, torch.bfloat16):
                logits = torch.randn(64, 151936, dtype=dtype, device="cuda")
                got_m, got_ls, got_v, got_i = row_logsumexp_topk(logits, k)
                # The (max, log_sum) pair is bitwise the non-fused kernel's.
                ref_m, ref_ls = row_logsumexp(logits)
                self.assertTrue(torch.equal(got_m, ref_m), (k, dtype))
                self.assertTrue(torch.equal(got_ls, ref_ls), (k, dtype))
                ref_v, ref_i = torch.topk(logits, k, dim=-1, sorted=True)
                self.assertTrue(torch.equal(got_v, ref_v.float()), (k, dtype))
                if dtype == torch.float32:
                    # fp32 randn is tie-free w.h.p.: indices match exactly.
                    self.assertTrue(torch.equal(got_i, ref_i), (k, dtype))
                else:
                    # bf16 has value ties; indices must point at the values.
                    self.assertTrue(
                        torch.equal(logits.gather(-1, got_i).float(), got_v),
                        (k, dtype),
                    )

    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
    def test_fused_topk_tie_break_is_lowest_index(self):
        from sglang.srt.layers.logsumexp import row_logsumexp_topk

        logits = torch.zeros(1, 1000, device="cuda")
        logits[0, [7, 3, 500]] = 5.0
        _, _, _, got_i = row_logsumexp_topk(logits, 3)
        self.assertEqual(got_i[0].tolist(), [3, 7, 500])

    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
    def test_fused_topk_inf_rows(self):
        from sglang.srt.layers.logsumexp import row_logsumexp_topk

        logits = torch.full((3, 1000), float("-inf"), device="cuda")
        logits[1, 3] = 2.5
        _, _, got_v, got_i = row_logsumexp_topk(logits.bfloat16(), 2)
        self.assertEqual(got_i[0].tolist(), [0, 1])
        self.assertEqual(got_i[1, 0].item(), 3)
        self.assertFalse(got_v.isnan().any().item())

    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
    def test_fast_path_end_to_end_on_cuda(self):
        # Exercises the fused-kernel integration inside _forward_by_chunk
        # (the CPU sweeps only cover the torch fallbacks), including the
        # k > FUSED_TOPK_MAX_K fallback.
        torch.manual_seed(0)
        proc = InputLogprobProcessor()
        for k_override in (None, 20):
            # k=20 exceeds FUSED_TOPK_MAX_K, exercising the torch fallback;
            # it needs a vocab that can supply 20 entries.
            batch = _build_batch([(4, 1), (6, 2), (3, 0)], torch.float32, vocab=64)
            pruned_states, sample_indices, lp_indices, t2s, metadata = batch
            if k_override is not None:
                metadata.top_logprobs_nums = [k_override] * 3
            batch = (
                pruned_states.cuda(),
                sample_indices.cuda(),
                lp_indices.cuda(),
                t2s,
                metadata,
            )
            metadata.extend_input_logprob_token_ids_gpu = (
                metadata.extend_input_logprob_token_ids_gpu.cuda()
            )
            for chunk_size in (None, 2, 5):
                ref, _ = _run(proc, batch, False, chunk_size)
                got, _ = _run(proc, batch, True, chunk_size)
                label = f"k_override={k_override} chunk={chunk_size}"
                self.assertEqual(ref.top_logprobs_idx, got.top_logprobs_idx, label)
                _assert_nested_close(
                    self, ref.top_logprobs_val, got.top_logprobs_val, label, 1e-5, 1e-5
                )
                torch.testing.assert_close(
                    ref.token_logprobs, got.token_logprobs, rtol=1e-5, atol=1e-5
                )

    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
    def test_row_logsumexp_kernel_matches_reference(self):
        from sglang.srt.layers.logsumexp import row_logsumexp

        torch.manual_seed(0)
        for rows, cols in ((0, 128), (3, 0), (1, 1), (7, 1000), (64, 151936)):
            for dtype in (torch.bfloat16, torch.float32):
                logits = torch.randn(rows, cols, dtype=dtype, device="cuda") * 8
                got_max, got_log_sum = row_logsumexp(logits)
                self.assertEqual(got_max.dtype, torch.float32)
                self.assertEqual(got_log_sum.dtype, torch.float32)
                if not cols:
                    self.assertTrue((got_max == float("-inf")).all())
                    self.assertTrue((got_log_sum == 0).all())
                    continue
                self.assertTrue(torch.equal(got_max, logits.float().amax(-1)))
                ref_log_sum = torch.logsumexp(
                    logits.double() - got_max.double()[:, None], dim=-1
                ).float()
                torch.testing.assert_close(
                    ref_log_sum,
                    got_log_sum,
                    rtol=1e-4,
                    atol=1e-4,
                    msg=f"{rows}x{cols} {dtype}",
                )
        # Rows dominated by -inf (masked-vocab shapes) must stay nan-free.
        logits = torch.full((4, 1000), float("-inf"), device="cuda")
        logits[1, 3] = 2.5
        logits[2, :] = torch.randn(1000, device="cuda")
        got_max, got_log_sum = row_logsumexp(logits.bfloat16())
        self.assertEqual(got_max[0].item(), float("-inf"))
        self.assertAlmostEqual((got_max[1] + got_log_sum[1]).item(), 2.5, delta=1e-2)
        self.assertFalse(got_max.isnan().any() or got_log_sum.isnan().any())
        # Non-contiguous input (sliced rows) exercises the stride args.
        base = torch.randn(8, 512, device="cuda", dtype=torch.bfloat16)
        view = base[::2]
        got_max, got_log_sum = row_logsumexp(view)
        torch.testing.assert_close(
            torch.logsumexp(view.double(), dim=-1).float(),
            got_max + got_log_sum,
            rtol=1e-4,
            atol=1e-4,
        )


if __name__ == "__main__":
    unittest.main()
