//! World-model sample-loading benchmark — clawhdf5 vs the h5py counterpart.
//!
//! Reproduces the access pattern of `stable-worldmodel`'s HDF5 dataloader
//! (arXiv 2605.21800): a dataset of `(N, H, W, C)` uint8 observation frames,
//! read one frame at a time in shuffled (dataloader) order. That paper
//! reports generic HDF5 at 1,416–1,474 samples/s (vs Lance 4,815); this
//! measures clawhdf5 and h5py on the **same machine and file**, so the
//! comparison is hardware-controlled. Absolute numbers are not comparable to
//! the paper's (different box, smaller frames, no torch/transform) — only
//! clawhdf5-vs-h5py *here* is.
//!
//! clawhdf5 mmaps the file once and takes a zero-copy `&[u8]` over the
//! contiguous observation dataset; frame `i` is a subslice, and the OS pages
//! it in on access. Two modes, because fairness demands both:
//!   * default: sum the frame bytes through the zero-copy view — clawhdf5's
//!     real advantage, no per-frame allocation;
//!   * `--copy`: `to_vec()` each frame first, matching h5py's unavoidable
//!     per-frame numpy materialization, so the two do equal work.
//!
//! Usage: `... --example worldmodel_sampling -- <file.h5> [passes] [--copy]`

use std::hint::black_box;
use std::time::Instant;

use clawhdf5::MmapFile;

fn main() {
    let args: Vec<String> = std::env::args().collect();
    let path = args
        .get(1)
        .expect("usage: worldmodel_sampling <file.h5> [passes] [--copy]");
    let passes: usize = args.get(2).and_then(|s| s.parse().ok()).unwrap_or(5);
    let copy = args.iter().any(|a| a == "--copy");

    let file = MmapFile::open(path).expect("open");
    let ds = file.dataset("observation").expect("observation dataset");
    let shape = ds.shape().expect("shape");
    let n = shape[0] as usize;
    let frame_bytes: usize = shape[1..].iter().map(|&d| d as usize).product();
    let raw = ds
        .read_raw_slice()
        .expect("read_raw_slice")
        .expect("contiguous zero-copy slice");
    assert_eq!(raw.len(), n * frame_bytes, "unexpected dataset size");

    let order = shuffled(n);

    let touch = |slice: &[u8]| -> u64 {
        if copy {
            let owned = slice.to_vec();
            owned.iter().map(|&b| u64::from(b)).sum()
        } else {
            slice.iter().map(|&b| u64::from(b)).sum()
        }
    };

    // Warm one pass (page-in), then time.
    let mut sink = 0u64;
    for &i in &order {
        sink = sink.wrapping_add(touch(&raw[i * frame_bytes..(i + 1) * frame_bytes]));
    }
    black_box(sink);

    let t0 = Instant::now();
    let mut sink = 0u64;
    for _ in 0..passes {
        for &i in &order {
            sink = sink.wrapping_add(touch(&raw[i * frame_bytes..(i + 1) * frame_bytes]));
        }
    }
    black_box(sink);
    let elapsed = t0.elapsed().as_secs_f64();

    let total = (n * passes) as f64;
    let mode = if copy {
        "materialized copy"
    } else {
        "zero-copy view"
    };
    println!("clawhdf5 ({mode}): {n} frames x {passes} passes in {elapsed:.3}s");
    println!("clawhdf5 ({mode}): {:.0} samples/sec", total / elapsed);
}

fn shuffled(n: usize) -> Vec<usize> {
    let mut v: Vec<usize> = (0..n).collect();
    let mut state: u64 = 0x9E37_79B9_7F4A_7C15;
    for i in (1..n).rev() {
        state = state
            .wrapping_mul(6364136223846793005)
            .wrapping_add(1442695040888963407);
        let j = (state >> 33) as usize % (i + 1);
        v.swap(i, j);
    }
    v
}