#!/usr/bin/env python3
"""Generate a world-model-shaped dataset: N frames of HxWxC uint8 observations,
contiguous (N,H,W,C), matching stable-worldmodel's per-frame sample-loading
access pattern. Also emits ep_len/ep_offset like their format."""
import sys, time, numpy as np, h5py

path = sys.argv[1]
N = int(sys.argv[2]) if len(sys.argv) > 2 else 20000
H = W = 64
C = 3
rng = np.random.default_rng(0)
t0 = time.perf_counter()
with h5py.File(path, "w", libver="latest") as f:
    # Contiguous (N,H,W,C) uint8 — the fair, both-APIs-support-it layout.
    obs = f.create_dataset("observation", shape=(N, H, W, C), dtype=np.uint8)
    # Write in blocks to bound memory.
    B = 2000
    for i in range(0, N, B):
        n = min(B, N - i)
        obs[i:i+n] = rng.integers(0, 256, size=(n, H, W, C), dtype=np.uint8)
    # Episode metadata like their format: 100-step episodes.
    ep = 100
    n_ep = N // ep
    f.create_dataset("ep_len", data=np.full(n_ep, ep, dtype=np.int32))
    f.create_dataset("ep_offset", data=(np.arange(n_ep) * ep).astype(np.int64))
print(f"wrote {N} frames {H}x{W}x{C} to {path} in {time.perf_counter()-t0:.1f}s "
      f"({N*H*W*C/1e6:.0f} MB)")