"""Check the Triton Philox against TileLang's `T.rng_rand()` stream.

The reference is not derived from documentation -- it is dumped from the real kernel
(`scripts/h20_rng_dump.py`), because `curand_init`'s seeding and subsequence layout are not
documented. This script re-dumps and compares, so it stays honest if the container's CUDA version
changes the seeding.

Both `subsequence=0` and a non-zero one are covered on purpose: the subsequence lands in the third
counter word, which is invisible when it is zero.

Exit code is non-zero on any mismatch, so it can gate the sampler port.
"""

from __future__ import annotations

import argparse
import sys

import torch

from deepinfer.model_executor.layers.tilelang_ops.triton.philox import philox_stream_tl

sys.path.insert(0, "scripts")
from h20_rng_dump import rng_dump_tl  # noqa: E402  (path set above so the script is a library here)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--n", type=int, default=16)
    parser.add_argument("--seeds", default="0,1,42,12345,18446744073709551615")
    parser.add_argument("--subseqs", default="0,1,2048,4097")
    args = parser.parse_args()

    n = args.n
    seeds = [int(x) for x in args.seeds.split(",")]
    subseqs = [int(x) for x in args.subseqs.split(",")]

    reference = rng_dump_tl(n)
    candidate = philox_stream_tl()

    failures = 0
    checked = 0
    for seed in seeds:
        for seq in subseqs:
            seed_t = torch.tensor([seed], dtype=torch.uint64, device="cuda")
            seq_t = torch.tensor([seq], dtype=torch.uint64, device="cuda")

            want = torch.zeros(n, dtype=torch.uint32, device="cuda")
            reference(seed_t, seq_t, want)

            got = torch.zeros(n, dtype=torch.uint32, device="cuda")
            candidate(seed_t, seq_t, got)

            checked += 1
            if not torch.equal(want, got):
                failures += 1
                want_cpu = want.cpu().tolist()
                got_cpu = got.cpu().tolist()
                first = next(i for i, (a, b) in enumerate(zip(want_cpu, got_cpu, strict=True)) if a != b)
                print(f"MISMATCH seed={seed} seq={seq}: first diff at {first}, want {want_cpu[first]} got {got_cpu[first]}")

    print(f"\n{checked - failures}/{checked} streams bit-identical")
    if failures:
        raise SystemExit(1)


if __name__ == "__main__":
    main()