//! A classical sampler for the 2022 time-domain Gaussian boson sampler's
//! instances (Madsen et al., Nature 606, 75 (2022)), and its samples scored
//! beside the device's against their exact ground truth, a signed receipt per
//! run.
//!
//! ```text
//! cargo run --release --example gbs_sampler --features quantum_gbs,quantum_engine,quantum_receipt -- DATA SET sample [options]
//! cargo run --release --example gbs_sampler --features quantum_gbs,quantum_engine,quantum_receipt -- DATA SET score [options]
//! ```
//!
//! `DATA` holds the data set as `wai/tools/quantum-gbs/fetch.sh` lays it out;
//! every file read is held to its pinned SHA-256 (`gbs_loops/common.rs`).
//!
//! **`sample`** splits the instance's lossy state (squeezed vacua `r.npy`
//! through the published `T`) into a pure part and classical noise, holds the
//! pure part as a matrix-product state of bond dimension `χ`, and draws
//! samples (`quantum_gbs_mps`). Options: `--chi C` (64), `--cutoff D` (4, the
//! pure part's photons per mode), `--floor F` (1e-12), `--shots K` (10,000),
//! `--seed S` (1), `--part P` (shots between progress lines, 2,000),
//! `--out PATH` (the samples, as the device's: `.npy`, `shots × 1 × modes`,
//! 8-bit), `--receipt PATH`, `--threads T`, `--host LABEL`.
//!
//! **`score`** counts the device's samples and each sampler's
//! (`--sampler A.npy,B.npy` with `--label a,b`; each file's SHA-256 is
//! recorded) and compares each with the exact ground truth, as the two
//! papers score: the total photon number's distribution (total variation
//! distance), the modes' means, every pair's second-order cumulant and every
//! triple's third-order cumulant within `--window W` (43) bins (slope,
//! correlation and distance to the exact values, and the share of the
//! distance sampling alone explains; for the pairs, also the correlation and
//! squared distance net of sampling noise, from products of distinct batches
//! only, with jackknife standard errors), and, with `--scores A-B` and
//! `--per-n K`, the cross entropy of the first `K` samples with `N` photons
//! for `N` in `A..=B` (a hafnian each). A sampler with fewer samples than
//! the device is also compared with the device's first samples of the same
//! count. Options: `--batches B` (100), `--triple-batches B3` (20),
//! `--replicas R` (20), `--receipt PATH`, `--threads T`, `--host LABEL`.

#![allow(clippy::needless_range_loop, clippy::too_many_lines)]

#[path = "gbs_loops/common.rs"]
mod common;

use std::path::{Path, PathBuf};
use std::time::Instant;

use serde_json::{Value, json};
use wai_quantum::quantum_gbs::{GaussianState, Input, Mat, PhotonTable};
use wai_quantum::quantum_gbs_mps::{MpsOptions, Sampler, photon_cumulant3, split};

fn main() {
    let a: Vec<String> = std::env::args().skip(1).collect();
    let usage = "usage: gbs_sampler DATA SET sample|score [options]";
    let data = PathBuf::from(a.first().expect(usage));
    let set = a.get(1).expect(usage).clone();
    let mode = a.get(2).expect(usage).clone();
    let modes = common::SETS.iter().find(|(s, _)| *s == set).unwrap_or_else(|| panic!("SET is one of fig2, fig3a, fig3b, fig4, figS15")).1;
    let opt = |k: &str| a.iter().position(|x| x == k).map(|i| a.get(i + 1).cloned().unwrap_or_else(|| panic!("{k} needs a value")));
    let receipt = opt("--receipt").map(PathBuf::from);
    let ((request, result), secs, energy) = common::metered(|| match mode.as_str() {
        "sample" => sample(&data, &set, modes, &opt),
        "score" => score(&data, &set, modes, &opt),
        _ => panic!("{usage}"),
    });
    let mut result = result;
    result["wall_seconds"] = json!(secs);
    result["energy_figure"] = energy;
    if let Some(h) = opt("--host") {
        result["host"] = json!(h);
    }
    result["peak_resident_bytes"] = peak_resident().map_or_else(|| json!({ "unread": "read from /proc/self/status (VmHWM) on Linux only; on macOS the run's log holds /usr/bin/time -l" }), |b| json!(b));
    let doc = common::seal(request, result, receipt.as_deref());
    println!("wall {secs:.1} s · energy {}", doc["result"]["energy_figure"].get("joules_micro").map_or_else(|| format!("unmetered ({})", doc["result"]["energy_figure"]["unmetered"].as_str().unwrap_or("")), |j| format!("{:.1} J", j.as_u64().unwrap_or(0) as f64 / 1e6)));
    if let Some(p) = &receipt {
        println!("receipt {} (seal verifies: {})", p.display(), common::verifies(&doc));
    }
}

/// The process's peak resident memory, from `/proc/self/status` (Linux).
fn peak_resident() -> Option<u64> {
    let s = std::fs::read_to_string("/proc/self/status").ok()?;
    let line = s.lines().find(|l| l.starts_with("VmHWM:"))?;
    line.split_whitespace().nth(1)?.parse::<u64>().ok().map(|kb| kb * 1024)
}

fn threads(opt: &dyn Fn(&str) -> Option<String>) -> usize {
    opt("--threads").map_or_else(|| std::thread::available_parallelism().map_or(1, |t| t.get()), |v| v.parse().expect("--threads: a number"))
}

/// The instance: the published squeezing (its first `m` values) and transfer matrix.
fn instance(root: &Path, set: &str, m: usize) -> (Vec<f64>, Mat) {
    let path = |f: &str| format!("{set}/{f}");
    let t = common::complex_matrix(&common::read_pinned(root, &path("T.npy")));
    let r = common::real_vector(&common::read_pinned(root, &path("r.npy")));
    assert_eq!(t.n, m, "the transfer matrix has the set's modes");
    (r[..m].to_vec(), t)
}

fn ground_truth(r: &[f64], t: &Mat) -> GaussianState {
    GaussianState::through(&r.iter().map(|&x| Input::Squeezed(x)).collect::<Vec<_>>(), t).expect("the ground truth")
}

fn sources(set: &str, files: &[&str]) -> Vec<Value> {
    let mut v: Vec<Value> = files.iter().map(|f| common::source(&format!("{set}/{f}"))).collect();
    v.push(json!({ "repository": common::REPOSITORY, "commit": common::COMMIT, "path": "LICENSE", "sha256": common::LICENSE_SHA256, "licence": "Apache-2.0" }));
    v
}

const EXPERIMENT: &str = "time-domain photonic Gaussian boson sampling: Madsen et al., Nature 606, 75 (2022), doi:10.1038/s41586-022-04725-x";
const METHOD: &str = "Oh, Liu, Alexeev, Fefferman and Jiang, Nature Physics 20, 1461 (2024), arXiv:2306.03709, written from the paper: the lossy state split into a pure part and a random displacement, the pure part a matrix-product state over the time-ordered modes, a displacement drawn per sample";

/// The samples as the device's are published: `.npy`, `|i1`, `shots × 1 × modes`, C order.
fn npy_bytes(rows: &[u8], shots: usize, modes: usize) -> Vec<u8> {
    let mut header = format!("{{'descr': '|i1', 'fortran_order': False, 'shape': ({shots}, 1, {modes}), }}");
    while (10 + header.len() + 1) % 64 != 0 {
        header.push(' ');
    }
    header.push('\n');
    let mut out = Vec::with_capacity(10 + header.len() + rows.len());
    out.extend_from_slice(b"\x93NUMPY\x01\x00");
    out.extend_from_slice(&(header.len() as u16).to_le_bytes());
    out.extend_from_slice(header.as_bytes());
    out.extend_from_slice(rows);
    out
}

fn sample(root: &Path, set: &str, m: usize, opt: &dyn Fn(&str) -> Option<String>) -> (Value, Value) {
    let num = |k: &str, d: usize| opt(k).map_or(d, |v| v.parse().unwrap_or_else(|_| panic!("{k}: a number")));
    let chi = num("--chi", 64);
    let cutoff = num("--cutoff", 4);
    let shots = num("--shots", 10_000);
    let part = num("--part", 2000).max(1);
    let seed = opt("--seed").map_or(1, |v| v.parse::<u64>().expect("--seed: a number"));
    let floor = opt("--floor").map_or(1e-12, |v| v.parse::<f64>().expect("--floor: a number"));
    let th = threads(opt);
    let out = opt("--out").map(PathBuf::from);
    let t0 = Instant::now();
    let (r, t) = instance(root, set, m);
    let gt = ground_truth(&r, &t);
    let (mean_gt, _) = gt.photon_moments();
    let mu_gt: f64 = mean_gt.iter().sum();
    let read_secs = t0.elapsed().as_secs_f64();
    let t1 = Instant::now();
    let sp = split(&r, &t).expect("the split");
    let split_secs = t1.elapsed().as_secs_f64();
    println!(
        "{set}: {m} modes, exact mean photons {mu_gt:.4}; split in {split_secs:.1} s: common transmission {:.4}, pure part {:.4} photons, noise {:.4}",
        sp.common_transmission,
        sp.pure_photons(),
        sp.noise_photons()
    );
    let t2 = Instant::now();
    let opts = MpsOptions { chi, cutoff, floor, threads: th };
    let sampler = Sampler::new(&sp, &opts).expect("the sampler");
    let build_secs = t2.elapsed().as_secs_f64();
    let mps = &sampler.mps;
    let center = m / 2;
    let max_bond = *mps.bonds.iter().max().unwrap();
    let min_kept = mps.kept.iter().copied().fold(1.0, f64::min);
    let min_site = mps.site_norm.iter().copied().fold(1.0, f64::min);
    let site_product: f64 = mps.site_norm.iter().product();
    let bytes = mps.entries() * 16;
    println!(
        "matrix-product state in {build_secs:.1} s: χ ≤ {max_bond} (centre {}), Schmidt weight kept ≥ {min_kept:.6} (centre {:.6}), site norms ≥ {min_site:.6} (product {site_product:.4}), {:.1} MB of tensors, noise rank {}",
        mps.bonds[center],
        mps.kept[center],
        bytes as f64 / 1e6,
        sampler.noise_rank()
    );
    let t3 = Instant::now();
    let mut rows = Vec::with_capacity(shots * m);
    let mut first = 0usize;
    while first < shots {
        let k = part.min(shots - first);
        rows.extend_from_slice(&sampler.sample(first as u64, k, seed, th));
        first += k;
        let el = t3.elapsed().as_secs_f64();
        println!("  {first}/{shots} samples, {el:.1} s, {:.4} s per sample", el / first as f64);
    }
    let sample_secs = t3.elapsed().as_secs_f64();
    assert!(rows.iter().all(|&x| x <= 127), "every count fits the device's 8-bit signed format");
    let mean_total = rows.iter().map(|&x| x as f64).sum::<f64>() / shots as f64;
    let file = npy_bytes(&rows, shots, m);
    let sha = common::hex(&common::sha256(&file));
    if let Some(p) = &out {
        std::fs::write(p, &file).unwrap_or_else(|e| panic!("writing {}: {e}", p.display()));
        println!("{shots} samples written to {} (sha256 {sha}), mean photons {mean_total:.4}", p.display());
    }
    let request = json!({
        "experiment": EXPERIMENT, "set": set, "modes": m,
        "source": sources(set, &["T.npy", "r.npy"]),
        "state": "squeezed vacua with r.npy's first m values through the published T (the paper's (r, T)), as quantum_gbs::GaussianState::through builds it",
        "method": METHOD,
        "split": "the largest loss common to every path (σ_max(T)²) moved to the inputs and taken out of each squeezed vacuum exactly; the rest of the state split by its Williamson form. A feasible split, not the paper's least-photon optimum (a semidefinite program)",
        "settings": { "chi": chi, "cutoff": cutoff, "floor": floor, "shots": shots, "seed": seed, "threads": th },
    });
    let result = json!({
        "set": set,
        "exact_mean_photons": mu_gt,
        "split": { "common_transmission": sp.common_transmission, "pure_photons": sp.pure_photons(), "noise_photons": sp.noise_photons() },
        "mps": {
            "bond_max": max_bond, "bond_centre": mps.bonds[center], "bonds": mps.bonds,
            "schmidt_weight_kept_min": min_kept, "schmidt_weight_kept_centre": mps.kept[center],
            "discarded_at_centre": 1.0 - mps.kept[center],
            "site_norm_min": min_site, "site_norm_product": site_product,
            "tensor_bytes": bytes, "noise_rank": sampler.noise_rank(),
        },
        "samples": { "shots": shots, "bytes": file.len(), "sha256": sha, "mean_photons": mean_total, "format": "npy |i1, shots × 1 × modes, C order" },
        "phases_seconds": { "read_and_ground_truth": read_secs, "split": split_secs, "build": build_secs, "sample": sample_secs, "per_sample": sample_secs / shots as f64 },
    });
    (request, result)
}

// ---------------------------------------------------------------------------
// Scoring
// ---------------------------------------------------------------------------

/// Run `f(i)` for `i in 0..n` on `threads` threads; results in index order.
fn parallel<T: Send>(n: usize, threads: usize, f: impl Fn(usize) -> T + Sync) -> Vec<T> {
    let mut slots: Vec<Option<T>> = (0..n).map(|_| None).collect();
    let threads = threads.clamp(1, n.max(1));
    let chunk = n.div_ceil(threads).max(1);
    std::thread::scope(|s| {
        for (t, part) in slots.chunks_mut(chunk).enumerate() {
            let f = &f;
            s.spawn(move || {
                for (k, slot) in part.iter_mut().enumerate() {
                    *slot = Some(f(t * chunk + k));
                }
            });
        }
    });
    slots.into_iter().map(|x| x.expect("every slot filled")).collect()
}

/// One batch of samples, counted: totals, first and second moments, and third
/// moments of the triples within the window (`i < j < l ≤ i + w`).
struct Batch {
    shots: u64,
    /// Entries above 127 (negative in the device's signed 8-bit format): counted, and read as no photons, as gbs_loops does.
    negative: u64,
    hist: Vec<u64>,
    s1: Vec<u64>,
    s2: Vec<u64>,
    s3: Vec<u64>,
}

fn count(rows: &[u8], m: usize, w: usize) -> Batch {
    let mut b = Batch { shots: 0, negative: 0, hist: Vec::new(), s1: vec![0; m], s2: vec![0; m * m], s3: vec![0; m * w * w] };
    let mut nz: Vec<(usize, u64)> = Vec::with_capacity(m);
    for row in rows.chunks(m) {
        b.shots += 1;
        nz.clear();
        let mut total = 0usize;
        for (i, &v) in row.iter().enumerate() {
            if v > 127 {
                b.negative += 1;
                continue;
            }
            if v > 0 {
                nz.push((i, u64::from(v)));
                total += v as usize;
            }
        }
        if b.hist.len() <= total {
            b.hist.resize(total + 1, 0);
        }
        b.hist[total] += 1;
        for (x, &(i, ni)) in nz.iter().enumerate() {
            b.s1[i] += ni;
            for (y, &(j, nj)) in nz.iter().enumerate().skip(x) {
                b.s2[i * m + j] += ni * nj;
                if y == x || j > i + w {
                    continue;
                }
                for &(l, nl) in &nz[y + 1..] {
                    if l > i + w {
                        break;
                    }
                    b.s3[(i * w + (j - i - 1)) * w + (l - i - 1)] += ni * nj * nl;
                }
            }
        }
    }
    b
}

fn merge(parts: &[Batch], m: usize, w: usize) -> Batch {
    let mut all = Batch { shots: 0, negative: 0, hist: Vec::new(), s1: vec![0; m], s2: vec![0; m * m], s3: vec![0; m * w * w] };
    for b in parts {
        all.shots += b.shots;
        all.negative += b.negative;
        if all.hist.len() < b.hist.len() {
            all.hist.resize(b.hist.len(), 0);
        }
        for (x, y) in all.hist.iter_mut().zip(&b.hist) {
            *x += y;
        }
        for (x, y) in all.s1.iter_mut().zip(&b.s1) {
            *x += y;
        }
        for (x, y) in all.s2.iter_mut().zip(&b.s2) {
            *x += y;
        }
        for (x, y) in all.s3.iter_mut().zip(&b.s3) {
            *x += y;
        }
    }
    all
}

/// The batch's pair cumulants (`i < j`, row by row) and windowed triple cumulants.
fn cumulants(b: &Batch, pairs: &[(usize, usize)], triples: &[(usize, usize, usize)], m: usize, w: usize) -> (Vec<f64>, Vec<f64>) {
    let k = b.shots as f64;
    let e1 = |i: usize| b.s1[i] as f64 / k;
    let e2 = |i: usize, j: usize| b.s2[i * m + j] as f64 / k;
    let c2 = pairs.iter().map(|&(i, j)| e2(i, j) - e1(i) * e1(j)).collect();
    let c3 = triples
        .iter()
        .map(|&(i, j, l)| {
            let e3 = b.s3[(i * w + (j - i - 1)) * w + (l - i - 1)] as f64 / k;
            e3 - e2(i, j) * e1(l) - e2(i, l) * e1(j) - e2(j, l) * e1(i) + 2.0 * e1(i) * e1(j) * e1(l)
        })
        .collect();
    (c2, c3)
}

fn mean_se(xs: &[f64]) -> (f64, f64) {
    let k = xs.len() as f64;
    let m = xs.iter().sum::<f64>() / k;
    let v = xs.iter().map(|x| (x - m) * (x - m)).sum::<f64>() / (k - 1.0);
    (m, (v / k).sqrt())
}

/// How observed cumulants compare with the exact ones: slopes (through the
/// origin, and with an intercept), the correlation, the distance relative to
/// the exact values' norm, the part of it sampling explains (from the batch
/// standard errors), and the root mean square standardised difference.
fn agreement(obs: &[f64], se: &[f64], exact: &[f64]) -> Value {
    let k = obs.len() as f64;
    let (mo, mh) = (obs.iter().sum::<f64>() / k, exact.iter().sum::<f64>() / k);
    let (mut sxy, mut sxx, mut syy, mut soh, mut shh, mut d2, mut n2, mut z2, mut v2) = (0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0);
    for ((o, h), s) in obs.iter().zip(exact).zip(se) {
        sxy += (o - mo) * (h - mh);
        sxx += (o - mo) * (o - mo);
        syy += (h - mh) * (h - mh);
        soh += o * h;
        shh += h * h;
        d2 += (o - h) * (o - h);
        n2 += h * h;
        z2 += ((o - h) / s) * ((o - h) / s);
        v2 += s * s;
    }
    json!({
        "count": obs.len(),
        "slope_through_origin": soh / shh,
        "slope": sxy / syy,
        "correlation": sxy / (sxx * syy).sqrt(),
        "distance": (d2 / n2).sqrt(),
        "distance_absolute": d2.sqrt(),
        "distance_beyond_sampling": ((d2 - v2).max(0.0) / n2).sqrt(),
        "distance_from_sampling": (v2 / n2).sqrt(),
        "rms_z": (z2 / k).sqrt(),
    })
}

/// `all` less `part`: the totals, first and second moments of every sample but the part's.
fn without(all: &Batch, part: &Batch) -> Batch {
    let mut hist = all.hist.clone();
    for (x, y) in hist.iter_mut().zip(&part.hist) {
        *x -= y;
    }
    Batch {
        shots: all.shots - part.shots,
        negative: all.negative - part.negative,
        hist,
        s1: all.s1.iter().zip(&part.s1).map(|(a, b)| a - b).collect(),
        s2: all.s2.iter().zip(&part.s2).map(|(a, b)| a - b).collect(),
        s3: Vec::new(),
    }
}

/// The least-squares slope, Pearson's correlation and the two-norm distance of
/// observed cumulants against exact ones, as the 2024 paper scores them.
fn line_metrics(obs: &[f64], exact: &[f64]) -> (f64, f64, f64) {
    let k = obs.len() as f64;
    let (mo, mh) = (obs.iter().sum::<f64>() / k, exact.iter().sum::<f64>() / k);
    let (mut sxy, mut sxx, mut syy, mut d2) = (0.0, 0.0, 0.0, 0.0);
    for (o, h) in obs.iter().zip(exact) {
        sxy += (o - mo) * (h - mh);
        sxx += (o - mo) * (o - mo);
        syy += (h - mh) * (h - mh);
        d2 += (o - h) * (o - h);
    }
    (sxy / syy, sxy / (sxx * syy).sqrt(), d2.sqrt())
}

/// The pair cumulants' correlation with the exact ones and their squared
/// distance from them, net of sampling noise. A batch's own noise enters only
/// its products with itself, so every sum here runs over distinct batches
/// `a ≠ b`, each batch's cumulants made unbiased (`n/(n − 1)`): over them,
/// `Σ (c_a − h)·(c_b − h) / (B(B − 1))` estimates the squared distance
/// `‖E c − h‖²`, the same sum with each batch's cumulants less their mean
/// estimates the spread of `E c`, and the correlation is the covariance of the
/// batches' mean with `h` over the square root of that spread times `h`'s.
/// The batches are contiguous, so drift across a source's samples is counted
/// as noise too, which can only bring that source nearer. Each estimate has its
/// jackknife standard error over the batches, each left out in turn.
fn net_of_sampling(parts: &[Batch], per_batch: &[(Vec<f64>, Vec<f64>)], exact: &[f64]) -> Value {
    let p = exact.len();
    let hbar = exact.iter().sum::<f64>() / p as f64;
    let hc: Vec<f64> = exact.iter().map(|h| h - hbar).collect();
    let shh: f64 = hc.iter().map(|x| x * x).sum();
    let nh: f64 = exact.iter().map(|x| x * x).sum::<f64>().sqrt();
    // Each batch's unbiased cumulants less their mean over the pairs (u), and less the exact ones (d).
    let rows: Vec<(Vec<f64>, Vec<f64>)> = parts
        .iter()
        .zip(per_batch)
        .filter(|(b, _)| b.shots > 1)
        .map(|(b, (c, _))| {
            let k = b.shots as f64;
            let c: Vec<f64> = c.iter().map(|x| x * k / (k - 1.0)).collect();
            let cbar = c.iter().sum::<f64>() / p as f64;
            (c.iter().map(|x| x - cbar).collect(), c.iter().zip(exact).map(|(x, h)| x - h).collect())
        })
        .collect();
    let nb = rows.len();
    let mut us = vec![0.0; p];
    let mut ds = vec![0.0; p];
    for (u, d) in &rows {
        for q in 0..p {
            us[q] += u[q];
            ds[q] += d[q];
        }
    }
    let dot = |a: &[f64], b: &[f64]| a.iter().zip(b).map(|(x, y)| x * y).sum::<f64>();
    // Per batch: ‖u‖², ‖d‖², U·u, D·d and u·(h − h̄); and the same summed.
    let per: Vec<[f64; 5]> = rows.iter().map(|(u, d)| [dot(u, u), dot(d, d), dot(&us, u), dot(&ds, d), dot(u, &hc)]).collect();
    let (uu, dd, g) = (dot(&us, &us), dot(&ds, &ds), per.iter().map(|x| x[4]).sum::<f64>());
    let (quu, qdd) = (per.iter().map(|x| x[0]).sum::<f64>(), per.iter().map(|x| x[1]).sum::<f64>());
    let estimate = |k: f64, uu: f64, quu: f64, dd: f64, qdd: f64, g: f64| -> (f64, f64) {
        let pairs = k * (k - 1.0);
        let spread = (uu - quu) / pairs;
        ((g / k) / (spread * shh).sqrt(), (dd - qdd) / pairs)
    };
    let (r, d2) = estimate(nb as f64, uu, quu, dd, qdd, g);
    let loo: Vec<(f64, f64)> = per.iter().map(|x| estimate((nb - 1) as f64, uu - 2.0 * x[2] + x[0], quu - x[0], dd - 2.0 * x[3] + x[1], qdd - x[1], g - x[4])).collect();
    json!({
        "batches": nb,
        "correlation": r,
        "correlation_se": jackknife(&loo.iter().map(|x| x.0).collect::<Vec<_>>()),
        "distance_squared_absolute": d2,
        "distance_squared_absolute_se": jackknife(&loo.iter().map(|x| x.1).collect::<Vec<_>>()),
        "distance": d2.max(0.0).sqrt() / nh,
    })
}

/// The jackknife standard error of a statistic from its leave-one-out values.
fn jackknife(xs: &[f64]) -> f64 {
    let k = xs.len() as f64;
    let m = xs.iter().sum::<f64>() / k;
    ((k - 1.0) / k * xs.iter().map(|x| (x - m) * (x - m)).sum::<f64>()).sqrt()
}

/// The mean, variance and third central moment (the second and third cumulants) of a
/// distribution over photon numbers given as weights.
fn moments3(w: &[f64]) -> (f64, f64, f64) {
    let total: f64 = w.iter().sum();
    let mu = w.iter().enumerate().map(|(n, x)| n as f64 * x).sum::<f64>() / total;
    let (mut k2, mut k3) = (0.0, 0.0);
    for (n, x) in w.iter().enumerate() {
        let d = n as f64 - mu;
        k2 += d * d * x;
        k3 += d * d * d * x;
    }
    (mu, k2 / total, k3 / total)
}

fn tvd(p: &[f64], q: &[f64]) -> f64 {
    let n = p.len().max(q.len());
    (0..n).map(|i| (p.get(i).copied().unwrap_or(0.0) - q.get(i).copied().unwrap_or(0.0)).abs()).sum::<f64>() / 2.0
}

struct Exact {
    m: usize,
    w: usize,
    p_total: Vec<f64>,
    mean: Vec<f64>,
    var: Vec<f64>,
    pairs: Vec<(usize, usize)>,
    c2: Vec<f64>,
    triples: Vec<(usize, usize, usize)>,
    c3: Vec<f64>,
    table: PhotonTable,
}

/// Every statistic of one set of samples (rows of photon numbers), as a JSON
/// value, with the observed total photon histogram normalised.
fn stats(rows: &[u8], ex: &Exact, batches: usize, tbatches: usize, replicas: usize, th: usize) -> (Value, Vec<f64>) {
    let (m, w) = (ex.m, ex.w);
    let shots = rows.len() / m;
    let per = shots.div_ceil(batches);
    let parts: Vec<Batch> = parallel(batches, th, |b| {
        let lo = (b * per).min(shots);
        let hi = ((b + 1) * per).min(shots);
        count(&rows[lo * m..hi * m], m, w)
    });
    let all = merge(&parts, m, w);
    let kf = all.shots as f64;
    let h: Vec<f64> = all.hist.iter().map(|&x| x as f64 / kf).collect();
    let mean_total = all.hist.iter().enumerate().map(|(n, &c)| n as f64 * c as f64).sum::<f64>() / kf;
    let floor: Vec<f64> = parallel(replicas, th, |rep| {
        let mut rng = common::Rng(0x5eed_0000 + rep as u64);
        let c = common::draw(&ex.p_total, all.shots, &mut rng);
        tvd(&c.iter().map(|&x| x as f64 / kf).collect::<Vec<_>>(), &ex.p_total)
    });
    let (fm, fse) = mean_se(&floor);
    // The total photon number's second and third cumulants, with standard errors from the batches.
    let (_, k2, k3) = moments3(&h);
    let (_, k2x, k3x) = moments3(&ex.p_total);
    let per_k: Vec<(f64, f64, f64)> = parts.iter().filter(|b| b.shots > 1).map(|b| moments3(&b.hist.iter().map(|&x| x as f64).collect::<Vec<_>>())).collect();
    let k2se = mean_se(&per_k.iter().map(|x| x.1).collect::<Vec<_>>()).1;
    let k3se = mean_se(&per_k.iter().map(|x| x.2).collect::<Vec<_>>()).1;
    // First order.
    let means: Vec<f64> = all.s1.iter().map(|&x| x as f64 / kf).collect();
    let z1: Vec<f64> = (0..m).map(|i| (means[i] - ex.mean[i]) / (ex.var[i] / kf).sqrt()).collect();
    let rel: Vec<f64> = (0..m).map(|i| means[i] / ex.mean[i] - 1.0).collect();
    // Second and third order, standard errors from the batches (pairs) and from coarser batches (triples).
    let (c2, c3) = cumulants(&all, &ex.pairs, &ex.triples, m, w);
    let per_batch: Vec<(Vec<f64>, Vec<f64>)> = parallel(batches, th, |b| cumulants(&parts[b], &ex.pairs, &[], m, w));
    let se2: Vec<f64> = (0..ex.pairs.len()).map(|p| mean_se(&per_batch.iter().map(|x| x.0[p]).collect::<Vec<_>>()).1).collect();
    let group = batches.div_ceil(tbatches);
    let coarse: Vec<Batch> = (0..tbatches).map(|g| merge(&parts[(g * group).min(batches)..((g + 1) * group).min(batches)], m, w)).filter(|b| b.shots > 0).collect();
    let per3: Vec<Vec<f64>> = parallel(coarse.len(), th, |g| cumulants(&coarse[g], &[], &ex.triples, m, w).1);
    let se3: Vec<f64> = (0..ex.triples.len()).map(|t| mean_se(&per3.iter().map(|x| x[t]).collect::<Vec<_>>()).1).collect();
    // Each headline statistic's jackknife standard error over the batches, each left out in turn.
    let loo: Vec<[f64; 7]> = parallel(batches, th, |b| {
        let rest = without(&all, &parts[b]);
        let kr = rest.shots as f64;
        let hr: Vec<f64> = rest.hist.iter().map(|&x| x as f64 / kr).collect();
        let (mu, k2r, k3r) = moments3(&hr);
        let (c2r, _) = cumulants(&rest, &ex.pairs, &[], m, w);
        let (sl, co, di) = line_metrics(&c2r, &ex.c2);
        [tvd(&hr, &ex.p_total), mu, k3r / k3x - 1.0, sl, co, di, k2r / k2x - 1.0]
    });
    let jk = |i: usize| jackknife(&loo.iter().map(|x| x[i]).collect::<Vec<_>>());
    let net = net_of_sampling(&parts, &per_batch, &ex.c2);
    let v = json!({
        "samples": all.shots,
        "negative_entries": all.negative,
        "totals": { "mean": mean_total, "tvd_to_exact": tvd(&h, &ex.p_total),
            "second_cumulant": { "observed": k2, "exact": k2x, "se": k2se, "relative_deviation": k2 / k2x - 1.0 },
            "third_cumulant": { "observed": k3, "exact": k3x, "se": k3se, "relative_deviation": k3 / k3x - 1.0 }, "tvd_from_sampling_alone": { "replicas": replicas, "mean": fm, "sd": fse * (replicas as f64).sqrt() }, "max": all.hist.len() - 1 },
        "first_order": {
            "rms_z": (z1.iter().map(|z| z * z).sum::<f64>() / m as f64).sqrt(),
            "mean_relative_deviation": rel.iter().sum::<f64>() / m as f64,
            "max_abs_relative_deviation": rel.iter().fold(0.0f64, |a, x| a.max(x.abs())),
        },
        "second_order": agreement(&c2, &se2, &ex.c2),
        "second_order_net_of_sampling": net,
        "jackknife_se": { "batches": batches, "tvd_to_exact": jk(0), "mean": jk(1), "second_cumulant_relative_deviation": jk(6), "third_cumulant_relative_deviation": jk(2), "second_order_slope": jk(3), "second_order_correlation": jk(4), "second_order_distance_absolute": jk(5) },
        "third_order": agreement(&c3, &se3, &ex.c3),
    });
    (v, h)
}

/// The paper's per-`N` cross entropy over the first `per_n` samples with `N` photons.
fn cross_entropy(rows: &[u8], ex: &Exact, lo: usize, hi: usize, per_n: usize, th: usize) -> Vec<Value> {
    let m = ex.m;
    let mut by_n: Vec<Vec<Vec<usize>>> = vec![Vec::new(); hi + 1];
    for row in rows.chunks(m) {
        let n: usize = row.iter().map(|&v| v as usize).sum();
        if (lo..=hi).contains(&n) && by_n[n].len() < per_n {
            by_n[n].push(row.iter().map(|&v| v as usize).collect());
        }
    }
    let ln_binom = |n: usize, k: usize| -> f64 { (0..k).map(|i| ((n - i) as f64).ln() - ((i + 1) as f64).ln()).sum() };
    let mut out = Vec::new();
    for n in lo..=hi {
        let s = &by_n[n];
        if s.len() < 2 {
            continue;
        }
        let outer = if n >= 12 { 1 } else { th };
        let p: Vec<f64> = parallel(s.len(), outer, |i| ex.table.probability(&s[i]).expect("a pattern"));
        let ln_uniform = ex.p_total[n].ln() - ln_binom(n + m - 1, n);
        let xe: Vec<f64> = p.iter().map(|x| x.ln() - ln_uniform).collect();
        let (mean, se) = mean_se(&xe);
        out.push(json!({ "n": n, "samples": s.len(), "mean": mean, "se": se }));
    }
    out
}

fn score(root: &Path, set: &str, m: usize, opt: &dyn Fn(&str) -> Option<String>) -> (Value, Value) {
    let num = |k: &str, d: usize| opt(k).map_or(d, |v| v.parse().unwrap_or_else(|_| panic!("{k}: a number")));
    let batches = num("--batches", 100);
    let tbatches = num("--triple-batches", 20);
    let replicas = num("--replicas", 20);
    let w = num("--window", 43);
    let per_n = num("--per-n", 1000);
    let scores = opt("--scores").map(|s| {
        let (x, y) = s.split_once('-').expect("--scores A-B");
        (x.parse::<usize>().expect("A"), y.parse::<usize>().expect("B"))
    });
    let th = threads(opt);
    let files: Vec<PathBuf> = opt("--sampler").map(|s| s.split(',').map(PathBuf::from).collect()).unwrap_or_default();
    let labels: Vec<String> = opt("--label").map(|s| s.split(',').map(str::to_string).collect()).unwrap_or_else(|| (0..files.len()).map(|i| format!("sampler{i}")).collect());
    assert_eq!(files.len(), labels.len(), "one label per sampler file");
    let t0 = Instant::now();
    let (r, t) = instance(root, set, m);
    let device_bytes = common::read_pinned(root, &format!("{set}/samples.npy"));
    let (dshots, dm, drows) = common::sample_rows(&device_bytes);
    assert_eq!(dm, m);
    let gt = ground_truth(&r, &t);
    let (mean, cov) = gt.photon_moments();
    let mu: f64 = mean.iter().sum();
    let nmax = (8.0 * mu) as usize + 200;
    let p_total = gt.photon_number_distribution(nmax).expect("the exact total");
    let pairs: Vec<(usize, usize)> = (0..m).flat_map(|i| (i + 1..m).map(move |j| (i, j))).collect();
    let c2 = pairs.iter().map(|&(i, j)| cov[i * m + j]).collect();
    let triples: Vec<(usize, usize, usize)> = (0..m).flat_map(|i| (i + 1..=(i + w).min(m - 1)).flat_map(move |j| (j + 1..=(i + w).min(m - 1)).map(move |l| (i, j, l)))).collect();
    let c3: Vec<f64> = parallel(triples.len(), th, |k| {
        let (i, j, l) = triples[k];
        photon_cumulant3(&gt, i, j, l)
    });
    let table = gt.photon_table().expect("the photon table");
    let ex = Exact { m, w, p_total, var: (0..m).map(|i| cov[i * m + i]).collect(), mean, pairs, c2, triples, c3, table };
    let exact_secs = t0.elapsed().as_secs_f64();
    println!("{set}: {m} modes, exact mean {mu:.4}; {} pairs, {} triples within {w} bins; exact statistics in {exact_secs:.1} s", ex.pairs.len(), ex.triples.len());
    let mut sets: Vec<(String, Value, Vec<f64>)> = Vec::new();
    let mut sampler_meta = Vec::new();
    let report = |name: &str, v: &Value| {
        println!(
            "{name:>24}: {:>8} samples · total mean {:.3}, TVD {:.4} (sampling alone {:.4}) · means {:+.4} · pairs slope {:.4}/{:.4} corr {:.4} dist {:.4} abs {:.4} (beyond sampling {:.4}; net of sampling corr {:.4} ± {:.4}, squared distance {:.3e} ± {:.3e}) · triples slope {:.4} corr {:.4} dist {:.4} (beyond sampling {:.4})",
            v["samples"], v["totals"]["mean"].as_f64().unwrap(), v["totals"]["tvd_to_exact"].as_f64().unwrap(), v["totals"]["tvd_from_sampling_alone"]["mean"].as_f64().unwrap(),
            v["first_order"]["mean_relative_deviation"].as_f64().unwrap(),
            v["second_order"]["slope_through_origin"].as_f64().unwrap(), v["second_order"]["slope"].as_f64().unwrap(), v["second_order"]["correlation"].as_f64().unwrap(), v["second_order"]["distance"].as_f64().unwrap(), v["second_order"]["distance_absolute"].as_f64().unwrap(), v["second_order"]["distance_beyond_sampling"].as_f64().unwrap(),
            v["second_order_net_of_sampling"]["correlation"].as_f64().unwrap_or(f64::NAN), v["second_order_net_of_sampling"]["correlation_se"].as_f64().unwrap_or(f64::NAN),
            v["second_order_net_of_sampling"]["distance_squared_absolute"].as_f64().unwrap_or(f64::NAN), v["second_order_net_of_sampling"]["distance_squared_absolute_se"].as_f64().unwrap_or(f64::NAN),
            v["third_order"]["slope_through_origin"].as_f64().unwrap(), v["third_order"]["correlation"].as_f64().unwrap(), v["third_order"]["distance"].as_f64().unwrap(), v["third_order"]["distance_beyond_sampling"].as_f64().unwrap(),
        );
    };
    let (dv, dh) = stats(&drows, &ex, batches, tbatches, replicas, th);
    report("device", &dv);
    sets.push(("device".into(), dv, dh));
    let mut sampler_rows = Vec::new();
    for (f, label) in files.iter().zip(&labels) {
        let bytes = std::fs::read(f).unwrap_or_else(|e| panic!("reading {}: {e}", f.display()));
        let sha = common::hex(&common::sha256(&bytes));
        let (shots, sm, rows) = common::sample_rows(&bytes);
        assert_eq!(sm, m, "{}: the sampler's samples have the set's modes", f.display());
        sampler_meta.push(json!({ "label": label, "file": f.file_name().map(|x| x.to_string_lossy().to_string()), "bytes": bytes.len(), "sha256": sha, "shots": shots }));
        let (v, h) = stats(&rows, &ex, batches, tbatches, replicas, th);
        report(label, &v);
        if shots < dshots && !sets.iter().any(|s| s.0 == format!("device_first_{shots}")) {
            let (dv, dh) = stats(&drows[..shots * m], &ex, batches, tbatches, replicas, th);
            report(&format!("device_first_{shots}"), &dv);
            sets.push((format!("device_first_{shots}"), dv, dh));
        }
        sets.push((label.clone(), v, h));
        sampler_rows.push((label.clone(), rows.into_owned()));
    }
    // Each sampler's total photon distribution against the device's.
    let device_h = sets[0].2.clone();
    let mut by_name = serde_json::Map::new();
    for (name, mut v, h) in sets {
        if name != "device" {
            v["totals"]["tvd_to_device"] = json!(tvd(&h, &device_h));
        }
        by_name.insert(name, v);
    }
    let mut result = json!({ "set": set, "modes": m, "exact_mean_photons": mu, "pairs": ex.pairs.len(), "triples": ex.triples.len(), "window": w, "samples": by_name, "samplers": sampler_meta });
    if let Some((lo, hi)) = scores {
        let t4 = Instant::now();
        let mut xe = serde_json::Map::new();
        let d = cross_entropy(&drows, &ex, lo, hi, per_n, th);
        for row in &d {
            println!("cross entropy N = {}: device {:.4} ± {:.4} ({} samples)", row["n"], row["mean"].as_f64().unwrap(), row["se"].as_f64().unwrap(), row["samples"]);
        }
        xe.insert("device".into(), json!(d));
        for (label, rows) in &sampler_rows {
            let s = cross_entropy(rows, &ex, lo, hi, per_n, th);
            for row in &s {
                println!("cross entropy N = {}: {label} {:.4} ± {:.4} ({} samples)", row["n"], row["mean"].as_f64().unwrap(), row["se"].as_f64().unwrap(), row["samples"]);
            }
            xe.insert(label.clone(), json!(s));
        }
        result["cross_entropy"] = json!({ "per_n_cap": per_n, "selection": "the first samples in file order with each N", "by_set": xe, "seconds": t4.elapsed().as_secs_f64(),
            "logarithms": "the platform's f64 ln: these scores are not bit-identical across platforms, unlike the counts" });
    }
    result["phases_seconds"] = json!({ "exact": exact_secs, "total": t0.elapsed().as_secs_f64() });
    let request = json!({
        "experiment": EXPERIMENT, "set": set,
        "source": sources(set, &["T.npy", "r.npy", "samples.npy"]),
        "ground_truth": "squeezed vacua with r.npy's first m values through the published T (the paper's (r, T))",
        "samplers": "files written by gbs_sampler sample; each named here by its SHA-256",
        "statistics": {
            "totals": "the total photon number's histogram against quantum_gbs::GaussianState::photon_number_distribution; sampling alone: draws of the same count from the exact distribution, SplitMix64 seeds 0x5eed0000 + replica",
            "first_order": "each mode's mean against photon_moments",
            "second_order": "every pair i < j: the sample covariance of the photon numbers against photon_moments; standard errors from the spread across contiguous batches",
            "third_order": "every triple i < j < l ≤ i + window: the sample third cumulant against quantum_gbs_mps::photon_cumulant3; standard errors from coarser batches",
            "jackknife": "for the total's distance and mean, its second and third cumulants and the pair cumulants' slope, correlation and distance: the jackknife over the batches, each left out in turn",
            "second_order_net_of_sampling": "the pair cumulants' Pearson correlation with the exact ones and their squared distance ‖E c − h‖², each from products of distinct batches' unbiased cumulants only (Σ over a ≠ b, divided by B(B − 1)), so that no batch's own noise enters; distance √max(0, that)/‖h‖; standard errors by the jackknife over the batches. Drift across contiguous batches counts as noise",
            "metrics": "slope_through_origin Σoh/Σh², slope and correlation of the least-squares line and Pearson's r, distance ‖o − h‖/‖h‖, distance_absolute ‖o − h‖, distance_from_sampling √(Σ se²)/‖h‖, distance_beyond_sampling √max(0, ‖o − h‖² − Σ se²)/‖h‖",
            "cross_entropy": "ln(Pr(S) / (Pr(N)/C(N+M−1, N))) per sample, Pr(S) a hafnian (quantum_gbs::PhotonTable)",
        },
        "settings": { "batches": batches, "triple_batches": tbatches, "replicas": replicas, "window": w, "scores": scores.map(|(a, b)| json!([a, b])), "per_n": per_n, "threads": th },
    });
    (request, result)
}