import os
import torch
from deepinfer.model_executor.layers.tilelang_ops.neuron.random_sampling import (
    two_level_sample, quantize_weights, acc_slots_for, rng_first_two, rng_sequence)
from deepinfer.model_executor.layers.tilelang_ops.random_sampling import random_sampling
from deepinfer.model_executor.layers.tilelang_ops.sampler import prepare_verification_metadata, prepare_verification_metadata_ref

for V in [8, 100, 1024, 4032, 8192, 129280]:
    Q = 64
    p = torch.rand(Q, V) ** 8; p[:, ::3] = 0; p = p / p.sum(1, keepdim=True)
    seeds = torch.randint(-2**62, 2**62, (Q,), dtype=torch.int64)
    pos = torch.randint(-5, 10**9, (Q,), dtype=torch.int32)
    out = two_level_sample(p, seeds, pos, 1, V)
    # independent: inverse CDF over columns ordered by (col % AS, col // AS)
    AS = acc_slots_for(V)
    fx = quantize_weights(p).to(torch.int64)
    cols = torch.arange(V)
    order = torch.argsort((cols % AS) * (V + 1) + cols // AS)
    cdf = fx[:, order].cumsum(1)
    total = cdf[:, -1]
    lo, hi = rng_first_two(seeds, rng_sequence(pos, 1))
    r = torch.tensor([((int(h) << 32) | int(l)) % int(t) for h, l, t in zip(hi, lo, total)])
    exp = order[(cdf <= r[:, None]).sum(1)]
    assert torch.equal(out.long(), exp), V
    assert (p.gather(1, out.long()[:, None]) > 0).all()
print("two-level == slot-ordered inverse CDF: OK")
# strided probs rows
p = torch.softmax(torch.randn(16, 2048), -1); big = torch.zeros(16, 4096); big[:, :2048] = p
a = random_sampling(big[:, :2048], torch.arange(16).to(torch.uint64), torch.zeros(16, dtype=torch.int32))
b = random_sampling(p.contiguous(), torch.arange(16).to(torch.uint64), torch.zeros(16, dtype=torch.int32))
assert torch.equal(a, b); print("strided rows OK")
# nextn=0 metadata path (num_steps=0) and seqlen=1 rows
seqlen = torch.ones(37, dtype=torch.int32); ftp = torch.arange(37, dtype=torch.int32)
for got, ref in zip(prepare_verification_metadata(ftp, seqlen, 37, 0), prepare_verification_metadata_ref(ftp, seqlen, 37)):
    assert torch.equal(got, ref)
print("prepare_verification_metadata num_steps=0 OK")