#!/usr/bin/env python3
"""The concurrent_read workload with h5py, on the files concurrent_read wrote.

libhdf5 serialises every API call under one global lock, and h5py holds its
own global lock around every call as well, so h5py *threads* cannot decode in
parallel. h5py users scale with *processes* instead; ``--executor processes``
measures that (each worker opens the file itself).

The workload mirrors ``crates/clawhdf5-bench/src/bin/concurrent_read.rs``:

* ``distinct``: every dataset read in full once per repetition; worker ``t``
  of ``T`` reads datasets ``t, t + T, ...``.
* ``same``: ``--slabs`` random ``--slab`` x ``--slab`` hyperslabs of ``d00``
  (slab ``j`` to worker ``j % T``), offsets from the same splitmix64 stream.

Each worker times itself from a start barrier; a repetition spans the earliest
start to the latest finish (CLOCK_MONOTONIC, comparable across processes).
Threads share one ``h5py.File`` per repetition; process workers open the file
inside the timed region (a few ms against reads of many MiB).

Generate the files first with the Rust harness (it writes ``manifest.json``),
then, for example::

    python concurrent_read_h5py.py --dir DIR --executor threads --json h5py-threads.json
    python concurrent_read_h5py.py --dir DIR --executor processes --json h5py-procs.json
"""

import argparse
import json
import multiprocessing as mp
import os
import platform
import socket
import sys
import threading
import time

import h5py
import numpy as np

M64 = (1 << 64) - 1


def splitmix64(state):
    """Return (new_state, value); the same stream as the Rust harness."""
    state = (state + 0x9E3779B97F4A7C15) & M64
    z = state
    z = ((z ^ (z >> 30)) * 0xBF58476D1CE4E5B9) & M64
    z = ((z ^ (z >> 27)) * 0x94D049BB133111EB) & M64
    return state, z ^ (z >> 31)


def value(k, i):
    """Element i (row-major) of dataset k, exactly as concurrent_read writes it."""
    _, noise = splitmix64(i ^ (k << 40))
    return np.float32((((i >> 6) % 16384) + k) + (noise & 0xFF) / 256.0)


def slab_offsets(seed, count, rows, cols, slab):
    s = seed
    out = []
    for _ in range(count):
        s, r = splitmix64(s)
        s, c = splitmix64(s)
        out.append((r % (rows - slab + 1), c % (cols - slab + 1)))
    return out


def now():
    return time.clock_gettime(time.CLOCK_MONOTONIC)


def work(f, mode, t, threads, m, slabs, slab, verify):
    """Worker t's share of one repetition on an open h5py.File."""
    n = m["rows"] * m["cols"]
    if mode == "distinct":
        for k in range(t, m["datasets"], threads):
            got = f[f"d{k:02d}"][...]
            assert got.size == n
            if verify:
                flat = got.reshape(-1)
                for i in (0, n // 3, n - 1):
                    assert flat[i] == value(k, i), f"d{k:02d}[{i}]"
    else:
        ds = f["d00"]
        cols = m["cols"]
        for r, c in slabs[t::threads]:
            got = ds[r : r + slab, c : c + slab]
            assert got.shape == (slab, slab)
            if verify:
                assert got[0, 0] == value(0, r * cols + c)
                last = (r + slab - 1) * cols + c + slab - 1
                assert got[-1, -1] == value(0, last)


# ----- process workers ------------------------------------------------------

_barrier = None


def _init(barrier):
    global _barrier
    _barrier = barrier


def _proc_task(task):
    path, mode, t, threads, m, slabs, slab = task
    _barrier.wait()
    start = now()
    with h5py.File(path, "r") as f:
        work(f, mode, t, threads, m, slabs, slab, False)
    return start, now()


def _noop(_):
    return os.getpid()


def run_threads(path, mode, threads, m, slabs, slab):
    spans = [None] * threads
    barrier = threading.Barrier(threads)
    with h5py.File(path, "r") as f:

        def body(t):
            barrier.wait()
            start = now()
            work(f, mode, t, threads, m, slabs, slab, False)
            spans[t] = (start, now())

        ts = [threading.Thread(target=body, args=(t,)) for t in range(threads)]
        for th in ts:
            th.start()
        for th in ts:
            th.join()
    return max(e for _, e in spans) - min(s for s, _ in spans)


def run_processes(pool, path, mode, threads, m, slabs, slab):
    tasks = [(path, mode, t, threads, m, slabs, slab) for t in range(threads)]
    # One task per worker: each blocks in the barrier until all T have
    # started, so no worker can take a second task.
    spans = pool.map(_proc_task, tasks, chunksize=1)
    return max(e for _, e in spans) - min(s for s, _ in spans)


def warm(path):
    with open(path, "rb") as fh:
        while fh.read(1 << 24):
            pass


def evict(path):
    fd = os.open(path, os.O_RDONLY)
    try:
        os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED)
    finally:
        os.close(fd)


def main():
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("--dir", default="concurrent-read-data")
    ap.add_argument("--executor", choices=["threads", "processes"], default="threads")
    ap.add_argument("--threads", default="1,2,4,8,16")
    ap.add_argument("--reps", type=int, default=3)
    ap.add_argument("--slab", type=int, default=256)
    ap.add_argument("--slabs", type=int, default=1024)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--cold", action="store_true")
    ap.add_argument("--modes", default="distinct,same")
    ap.add_argument("--layouts", default="deflate,contiguous")
    ap.add_argument("--json")
    a = ap.parse_args()

    # The Rust harness pins this value (splitmix64_reference).
    assert splitmix64(42)[1] == 0xBDD732262FEB6E95, "splitmix64 port is wrong"

    try:
        with open(os.path.join(a.dir, "manifest.json")) as fh:
            m = json.load(fh)
    except FileNotFoundError:
        sys.exit(f"{a.dir}/manifest.json not found: generate the files with "
                 "`cargo run --release -p clawhdf5-bench --bin concurrent_read -- --dir ...` first")
    threads_list = [int(x) for x in a.threads.split(",")]
    modes = a.modes.split(",")
    layouts = a.layouts.split(",")
    if a.slab < 1 or a.slab > min(m["rows"], m["cols"]):
        sys.exit(f"--slab must be 1..={min(m['rows'], m['cols'])}")
    files = dict(m["files"])
    slabs = slab_offsets(a.seed, a.slabs, m["rows"], m["cols"], a.slab)
    dataset_bytes = m["rows"] * m["cols"] * 4
    tool = f"h5py-{a.executor}"

    ctx = mp.get_context("spawn")  # never fork a process holding HDF5 state
    pools = {}
    if a.executor == "processes":
        for t in threads_list:
            pool = ctx.Pool(t, initializer=_init, initargs=(ctx.Barrier(t),))
            pool.map(_noop, range(t))  # start the workers outside the timing
            pools[t] = pool

    rows = []
    print("| layout | mode | threads | MB/s | efficiency | median s |")
    print("|---|---|---:|---:|---:|---:|")
    try:
        for layout in layouts:
            path = os.path.join(a.dir, files[layout])
            if not a.cold:
                warm(path)
            for mode in modes:
                with h5py.File(path, "r") as f:  # untimed, checked pass
                    work(f, mode, 0, 1, m, slabs, a.slab, True)
                nbytes = (dataset_bytes * m["datasets"] if mode == "distinct"
                          else a.slab * a.slab * 4 * a.slabs)
                base = None
                for t in threads_list:
                    times = []
                    for _ in range(a.reps):
                        if a.cold:
                            evict(path)
                        if a.executor == "threads":
                            times.append(run_threads(path, mode, t, m, slabs, a.slab))
                        else:
                            times.append(run_processes(pools[t], path, mode, t, m, slabs, a.slab))
                    med = sorted(times)[len(times) // 2]
                    mb_s = nbytes / (1 << 20) / med
                    if t == 1:
                        base = mb_s
                    eff = mb_s / (t * base) if base else None
                    print(f"| {layout} | {mode} | {t} | {mb_s:.0f} | "
                          f"{'-' if eff is None else f'{eff:.2f}'} | {med:.4f} |")
                    rows.append({
                        "layout": layout, "mode": mode, "threads": t, "bytes": nbytes,
                        "times_s": times, "median_s": med, "mb_s": mb_s, "efficiency": eff,
                    })
    finally:
        for pool in pools.values():
            pool.terminate()

    if a.json:
        doc = {
            "tool": tool,
            "version": h5py.__version__,
            "hdf5_version": h5py.version.hdf5_version,
            "python": platform.python_version(),
            "host": socket.gethostname(),
            "cpus": os.cpu_count(),
            "unix_time": int(time.time()),
            "cache": ("cold (posix_fadvise DONTNEED before each repetition)"
                      if a.cold else "warm"),
            "decode_threads": 1,
            "params": {
                "datasets": m["datasets"], "rows": m["rows"], "cols": m["cols"],
                "chunk": m["chunk"], "deflate_level": m["deflate_level"],
                "mib": dataset_bytes // (1 << 20), "slab": a.slab, "slabs": a.slabs,
                "seed": a.seed, "reps": a.reps, "dir": a.dir,
            },
            "results": rows,
        }
        with open(a.json, "w") as fh:
            json.dump(doc, fh, indent=2)


if __name__ == "__main__":
    main()