"""Check the torch fallback for `deep_select.interface.topk` against a reference.

Covers the parts of the contract the indexer and sampler actually depend on:
`end` truncation (scores past it are uninitialized garbage), OOB padding when a row has
fewer than `topk` valid entries, `output_idx_offset`, in-place writes into a non-contiguous
view, and descending order for `sorted_index=True`.
"""

import torch

from deepinfer.h20_ops.deep_select_sm90 import topk


def ref_topk(
    scores: torch.Tensor,
    k: int,
    end: torch.Tensor | None = None,
    offset: int | None = None,
    fill: int = -1,
    val_fill: float = float("-inf"),
) -> tuple[torch.Tensor, torch.Tensor]:
    """Reference: consider only [0, end) per row, take k largest, pad, then shift."""
    rows = scores.shape[0]
    idx_rows, val_rows = [], []
    for r in range(rows):
        n = int(end[r]) if end is not None else scores.shape[1]
        row = scores[r, :n]
        kk = min(k, n)
        order = torch.argsort(row, descending=True)[:kk]
        vals = row[order]
        idxs = order.long()
        if offset is not None:
            idxs = idxs + int(offset[r])
        pad = k - kk
        if pad:
            idxs = torch.cat([idxs, torch.full((pad,), fill, dtype=idxs.dtype, device=scores.device)])
            vals = torch.cat([vals, torch.full((pad,), val_fill, dtype=vals.dtype, device=scores.device)])
        idx_rows.append(idxs)
        val_rows.append(vals)
    return torch.stack(val_rows), torch.stack(idx_rows)


def main() -> None:
    dev = "cuda"
    torch.manual_seed(0)
    rows, width, k = 6, 4096, 128

    scores = torch.randn(rows, width, device=dev, dtype=torch.float32)
    # Garbage past `end` that would win if the implementation ignored the bound.
    end = torch.tensor([4096, 4096, 200, 128, 3, 50], dtype=torch.int32, device=dev)
    for r, n in enumerate(end.tolist()):
        scores[r, n:] = 1e6

    vals, idxs = topk(scores, k, sorted=False, end=end, indices_type=torch.int32, return_value=True)
    # Default `idx_oob_fill_value` is INT_MAX, so the reference must pad with the same value.
    rv, ri = ref_topk(scores, k, end=end, fill=2147483647)
    assert idxs.dtype == torch.int32
    bad_rows = [r for r in range(rows) if not torch.equal(vals[r], rv[r]) or not torch.equal(idxs[r].long(), ri[r])]
    print(f"[1] end truncation + OOB pad: rows mismatched = {bad_rows}")
    assert not bad_rows
    # Rows shorter than topk must be padded with the fill value, not with garbage indices.
    short = (idxs[4] == 2147483647).sum().item()
    print(f"    row 4 (end=3) padded slots = {short} (expect {k - 3})")
    assert short == k - 3

    # Descending order for the valid prefix, which is what `sorted_index=True` promises.
    lens = torch.clamp(end.long(), max=k)
    for r in range(rows):
        n = int(lens[r])
        assert torch.equal(vals[r, :n], torch.sort(vals[r, :n], descending=True).values)
    print("[2] descending order holds")

    # In-place write into a non-contiguous view + per-row offset.
    storage = torch.full((rows, k * 2), -7, dtype=torch.int32, device=dev)
    view = storage[:, :k]
    offset = torch.arange(rows, dtype=torch.int32, device=dev) * 1000
    vals2, idxs2 = topk(
        scores,
        k,
        end=end,
        indices_type=torch.int32,
        output_idx=view,
        output_idx_offset=offset,
        idx_oob_fill_value=-1,
        return_value=False,
    )
    assert idxs2.data_ptr() == view.data_ptr()
    assert vals2 is None
    rv2, ri2 = ref_topk(scores, k, end=end, offset=offset, fill=-1)
    print(f"[3] in-place view: match = {torch.equal(view.long(), ri2)}")
    assert torch.equal(view.long(), ri2)
    # The fill value must not be shifted by the offset.
    assert int(view[4].min()) == -1

    # No `end` at all takes the plain path.
    v3, i3 = topk(scores, k, indices_type=torch.int64)
    rv3, ri3 = ref_topk(scores, k)
    print(f"[4] no end: values match = {torch.allclose(v3, rv3)}, indices match = {torch.equal(i3, ri3)}")
    assert torch.allclose(v3, rv3) and torch.equal(i3, ri3)

    # bf16 rows (the candidate-block path feeds bf16 block scores).
    bs = torch.randn(3, 512, device=dev, dtype=torch.bfloat16)
    vb, ib = topk(bs, 16, end=torch.tensor([512, 40, 8], dtype=torch.int32, device=dev), indices_type=torch.int32, return_value=False)
    padded = int((ib[2] == 2147483647).sum())
    print(f"[5] bf16: rows={tuple(ib.shape)} row2 padded={padded} (expect 8)")
    assert padded == 16 - 8

    print("PASS: deep_select torch fallback matches the documented contract")


if __name__ == "__main__":
    main()