//! HDF5 read-path measurement harness: full reads vs. hyperslab selections on
//! a chunked 2-D dataset, compressed and uncompressed, plus a contiguous one.
//!
//! The question it answers for every read-path change: does the cost of a
//! selection scale with the *selection*, or with the whole dataset?
//!
//! ```text
//! cargo run --release -p clawhdf5-bench --bin read_harness
//! cargo run --release -p clawhdf5-bench --bin read_harness -- --large   # 512 MB
//! ```

use std::time::{Duration, Instant};

use clawhdf5::{File, FileBuilder};
use clawhdf5_format::selection::Selection;

const CHUNK: u64 = 256;

struct Layout {
    name: &'static str,
    chunked: bool,
    deflate: bool,
}

const LAYOUTS: [Layout; 3] = [
    Layout {
        name: "chunked + deflate",
        chunked: true,
        deflate: true,
    },
    Layout {
        name: "chunked",
        chunked: true,
        deflate: false,
    },
    Layout {
        name: "contiguous",
        chunked: false,
        deflate: false,
    },
];

/// Smooth-ish, compressible data whose value encodes its position, so a read
/// can be verified exactly.
fn value(row: u64, col: u64) -> f64 {
    (row * 100_003 + col) as f64 * 0.5
}

fn write_file(path: &std::path::Path, rows: u64, cols: u64) {
    let data: Vec<f64> = (0..rows)
        .flat_map(|r| (0..cols).map(move |c| value(r, c)))
        .collect();
    let mut builder = FileBuilder::new();
    for (i, layout) in LAYOUTS.iter().enumerate() {
        let ds = builder.create_dataset(&format!("d{i}"));
        ds.with_f64_data(&data).with_shape(&[rows, cols]);
        if layout.chunked {
            ds.with_chunks(&[CHUNK, CHUNK]);
        }
        if layout.deflate {
            ds.with_deflate(4);
        }
    }
    builder.write(path).unwrap();
}

fn median(mut samples: Vec<Duration>) -> Duration {
    samples.sort();
    samples[samples.len() / 2]
}

fn time<T>(reps: usize, mut f: impl FnMut() -> T) -> Duration {
    median(
        (0..reps)
            .map(|_| {
                let t = Instant::now();
                std::hint::black_box(f());
                t.elapsed()
            })
            .collect(),
    )
}

fn slab(start: [u64; 2], count: [u64; 2]) -> Selection {
    Selection::Hyperslab {
        start: start.to_vec(),
        stride: vec![1, 1],
        count: count.to_vec(),
        block: vec![1, 1],
    }
}

fn main() {
    let large = std::env::args().any(|a| a == "--large");
    let (rows, cols) = if large { (8192, 8192) } else { (4096, 2048) };
    let total_mb = (rows * cols * 8) as f64 / (1 << 20) as f64;
    if cfg!(debug_assertions) {
        eprintln!("warning: debug build — numbers are meaningless. Use --release.");
    }

    let dir = tempfile::TempDir::new().unwrap();
    let path = dir.path().join("read_harness.h5");
    write_file(&path, rows, cols);
    let file_mb = std::fs::metadata(&path).unwrap().len() as f64 / (1 << 20) as f64;

    println!("## Read harness");
    println!(
        "\n{rows} x {cols} f64 ({total_mb:.0} MB per dataset), chunks {CHUNK} x {CHUNK}, file {file_mb:.0} MB\n"
    );

    // (label, selection, elements selected)
    let selections: Vec<(&str, Selection, u64)> = vec![
        (
            "64 x 64 window (1 chunk)",
            slab([300, 300], [64, 64]),
            64 * 64,
        ),
        (
            "512 x 512 window (4-9 chunks)",
            slab([1000, 700], [512, 512]),
            512 * 512,
        ),
        ("one row", slab([rows / 2, 0], [1, cols]), cols),
        ("one column", slab([0, cols / 2], [rows, 1]), rows),
    ];

    println!("| layout | read | selected | time ms | MB/s of selection | vs full read |");
    println!("|---|---|---:|---:|---:|---:|");
    for (i, layout) in LAYOUTS.iter().enumerate() {
        // Fresh handle per layout so one dataset's cached chunks don't help
        // (or evict) another's.
        let file = File::open(&path).unwrap();
        let ds = file.dataset(&format!("d{i}")).unwrap();

        let full_cold = time(1, || ds.read_f64().unwrap());
        let full = time(3, || ds.read_f64().unwrap());
        println!(
            "| {} | full (first) | {total_mb:.0} MB | {:.1} | {:.0} | |",
            layout.name,
            full_cold.as_secs_f64() * 1e3,
            total_mb / full_cold.as_secs_f64()
        );
        println!(
            "| {} | full (repeat) | {total_mb:.0} MB | {:.1} | {:.0} | 1.00x |",
            layout.name,
            full.as_secs_f64() * 1e3,
            total_mb / full.as_secs_f64()
        );

        for (label, selection, elements) in &selections {
            // A fresh handle again: measure the selection on its own, not
            // served from chunks the full read just cached.
            let file = File::open(&path).unwrap();
            let ds = file.dataset(&format!("d{i}")).unwrap();
            let got = ds.read_f64_selection(selection).unwrap();
            assert_eq!(got.len() as u64, *elements, "{label}");
            if let Selection::Hyperslab { start, .. } = selection {
                assert_eq!(got[0], value(start[0], start[1]), "{label}: wrong data");
            }
            let took = time(5, || {
                let file = File::open(&path).unwrap();
                let ds = file.dataset(&format!("d{i}")).unwrap();
                ds.read_f64_selection(selection).unwrap()
            });
            let mb = (*elements * 8) as f64 / (1 << 20) as f64;
            println!(
                "| {} | {label} | {:.2} MB | {:.2} | {:.0} | {:.3}x |",
                layout.name,
                mb,
                took.as_secs_f64() * 1e3,
                mb / took.as_secs_f64(),
                took.as_secs_f64() / full_cold.as_secs_f64()
            );
        }
    }
}