#!/usr/bin/env python3
"""h5py counterpart to worldmodel_sampling.rs — same file, same shuffled
per-frame access, same minimal touch (sum the frame bytes). Reports
samples/sec so the two sit side by side on one machine."""
import sys, time, numpy as np, h5py

path = sys.argv[1]
passes = int(sys.argv[2]) if len(sys.argv) > 2 else 5

def shuffled(n):
    v = list(range(n))
    state = 0x9E3779B97F4A7C15
    for i in range(n - 1, 0, -1):
        state = (state * 6364136223846793005 + 1442695040888963407) & 0xFFFFFFFFFFFFFFFF
        j = (state >> 33) % (i + 1)
        v[i], v[j] = v[j], v[i]
    return v

# swmr + a 256 MB chunk cache: exactly stable-worldmodel's HDF5Dataset._open_h5.
f = h5py.File(path, "r", swmr=True, rdcc_nbytes=256 * 1024 * 1024)
d = f["observation"]
n = d.shape[0]
order = shuffled(n)

# warm
sink = 0
for i in order:
    sink += int(d[i].sum())

t0 = time.perf_counter()
sink = 0
for _ in range(passes):
    for i in order:
        sink += int(d[i].sum())
elapsed = time.perf_counter() - t0
total = n * passes
print(f"h5py: {n} frames x {passes} passes = {total} reads in {elapsed:.3f}s")
print(f"h5py: {total/elapsed:.0f} samples/sec")