"""Dump TileLang's `T.rng_rand()` stream so a Triton implementation can be checked against it.

`T.rng_init(seed, seq, off, generator)` lowers to `curand_init(seed, seq, off, &state)` and
`T.rng_rand()` to `curand(&state)`, with the generator defaulting to `curandStatePhilox4_32_10_t`.
That is cuRAND's Philox4x32-10, whose mixing function is public but whose `curand_init` seeds and
fast-forwards the counter in a way that is not documented. So the stream is measured here rather
than derived, and `h20_philox_check.py` compares against it.

Prints, for each seed: the first `n` values, which is enough to pin both the counter progression and
the key schedule, since any mismatch in either shows up in the first few outputs.
"""

from __future__ import annotations

import argparse
import json

import tilelang
import tilelang.language as T
import torch


@tilelang.jit(
    pass_configs={
        tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
        tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
    },
)
def rng_dump_tl(n: int) -> object:
    @T.prim_func
    def _rng_dump(
        seed: T.Tensor((1,), T.uint64),
        seq: T.Tensor((1,), T.uint64),
        out: T.Tensor((n,), T.uint32),
    ) -> None:
        with T.Kernel(1, threads=32) as _:
            T.rng_init(seed[0], seq[0])
            for i in T.serial(n):
                # Every thread advances its own state, which is per-thread; only lane 0 writes, so
                # the recorded sequence is exactly that lane's.
                value = T.rng_rand()
                if T.get_thread_binding() == 0:
                    out[i] = value

    return _rng_dump


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

    n = args.n
    seeds = [int(x) for x in args.seeds.split(",")]
    kernel = rng_dump_tl(n)
    payload: dict[str, list[int]] = {}
    for seed in seeds:
        seed_t = torch.tensor([seed], dtype=torch.uint64, device="cuda")
        # Subsequence 0 is what the sampling kernel uses: `T.rng_init(seed, seq)` passes the row's
        # sample position as `seq`, so every distinct position is its own stream.
        for seq in (0, 1):
            out = torch.zeros(n, dtype=torch.uint32, device="cuda")
            kernel(seed_t, torch.tensor([seq], dtype=torch.uint64, device="cuda"), out)
            key = f"{seed}:{seq}"
            payload[key] = [int(x) for x in out.cpu()]
            print(f"seed={seed} seq={seq}: {payload[key]}", flush=True)

    if args.out:
        with open(args.out, "w") as f:
            json.dump(payload, f, indent=2)
        print(f"wrote {args.out}", flush=True)


if __name__ == "__main__":
    main()