//! The published samples of the 2022 time-domain Gaussian boson sampler
//! (Madsen et al., Nature 606, 75 (2022)) held to their ground truth, computed
//! exactly, with a signed receipt per run.
//!
//! ```text
//! cargo run --release --example gbs_loops --features quantum_gbs,quantum_engine,quantum_receipt -- DATA SET [options]
//! ```
//!
//! `DATA` holds the data set as its bucket lays it out (`DATA/fig2/T.npy`, …);
//! every file read is held to its pinned SHA-256 (`gbs_loops/common.rs`). `SET`
//! is one of `fig2` (16 modes), `fig3a` (216), `fig3b` (72), `fig4` (216) and
//! `figS15` (288). Options:
//!
//! - `--receipt PATH`: write the signed receipt there;
//! - `--counts PATH`: with `fig2`, write the observed count of every pattern
//!   of up to four photons (the fixture the tests re-run);
//! - `--max-n N`: with `fig2`, compare every pattern of up to `N` photons (6,
//!   as the paper does);
//! - `--scores A-B`, `--per-n K`: the paper's per-`N` scores for `N` in `A..=B`
//!   over the first `K` samples with `N` photons: the cross entropy against
//!   the ground truth (`fig3a`, as its Fig. 3a; 10,000 samples a point there)
//!   or the Bayesian log-average against each classical hypothesis (`fig3b`,
//!   as its Fig. 3b; 2,000 there);
//! - `--replicas R` (20), `--batches B` (100), `--threads T` (all).
//!
//! What it does, every step exact or counted:
//!
//! 1. rebuilds the instance's transfer matrix and unitary from its loop
//!    program (`quantum_gbs::Loops`, the paper's Methods eq. 10) and measures
//!    them against the published `T` and `U`;
//! 2. builds the ground truth, squeezed vacua `r` through `T`
//!    (`GaussianState::through`), and four classical hypotheses with the same
//!    mean photon number in every mode: thermal and squashed light through `T`,
//!    coherent light, and distinguishable squeezed light;
//! 3. counts the samples: the total photon number, every mode's mean, every
//!    pair's covariance (with standard errors from `B` batches), and, for the
//!    16-mode set, every pattern;
//! 4. compares: the total photon-number distribution with each hypothesis's,
//!    exact (`photon_number_distribution`); the first- and second-order
//!    cumulants with each hypothesis's, exact (`photon_moments`); for the
//!    16-mode set, every pattern's frequency with its exact probability
//!    (`PhotonTable`, a hafnian each), as fidelity and total variation
//!    distance within each photon number; and, for each, what the sample count
//!    alone would give, from `R` sets of draws from the exact distribution.

#![allow(clippy::needless_range_loop)]

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

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

use serde_json::{Value, json};
use wai_quantum::quantum_gbs::{GaussianState, Input, Mat, PhotonTable};

struct Args {
    data: PathBuf,
    set: String,
    receipt: Option<PathBuf>,
    counts: Option<PathBuf>,
    max_n: usize,
    scores: Option<(usize, usize)>,
    per_n: usize,
    replicas: usize,
    batches: usize,
    threads: usize,
}

fn main() {
    let a: Vec<String> = std::env::args().skip(1).collect();
    let usage = "usage: gbs_loops DATA SET [--receipt PATH] [--counts PATH] [--max-n N] [--scores A-B] [--per-n K] [--replicas R] [--batches B] [--threads T]";
    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 num = |k: &str, d: usize| opt(k).map_or(d, |v| v.parse().unwrap_or_else(|_| panic!("{k}: a number")));
    let set = a.get(1).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 args = Args {
        data: PathBuf::from(a.first().expect(usage)),
        receipt: opt("--receipt").map(PathBuf::from),
        counts: opt("--counts").map(PathBuf::from),
        max_n: num("--max-n", 6).min(15),
        scores: opt("--scores").map(|s| {
            let (x, y) = s.split_once('-').expect("--scores A-B");
            (x.parse().expect("A"), y.parse().expect("B"))
        }),
        per_n: num("--per-n", if set == "fig3b" { 2000 } else { 10_000 }),
        replicas: num("--replicas", 20),
        batches: num("--batches", 100),
        threads: num("--threads", std::thread::available_parallelism().map_or(1, |t| t.get())),
        set,
    };
    let ((request, result), secs, energy) = common::metered(|| run(&args, modes));
    let mut result = result;
    result["wall_seconds"] = json!(secs);
    result["energy_figure"] = energy;
    let doc = common::seal(request, result, args.receipt.as_deref());
    println!("wall {secs:.1} s · energy {}", doc["result"]["energy_figure"].get("joules_micro").map_or("unmetered".to_string(), |j| format!("{:.1} J", j.as_u64().unwrap_or(0) as f64 / 1e6)));
    if let Some(p) = &args.receipt {
        println!("receipt {} (seal verifies: {})", p.display(), common::verifies(&doc));
    }
}

/// `sinh² r`, the mean photon number of squeezed vacuum, from the engine.
fn nbar(r: f64) -> f64 {
    GaussianState::through(&[Input::Squeezed(r)], &Mat::identity(1)).expect("a squeezed mode").photon_moments().0[0]
}

/// 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.
struct Batch {
    shots: u64,
    hist: Vec<u64>,
    s1: Vec<u64>,
    s2: Vec<u64>,
    max_count: Vec<u8>,
    negative: u64,
    patterns: HashMap<u64, u64>,
}

fn count_batch(rows: &[u8], modes: usize, max_n: Option<usize>) -> Batch {
    let mut b = Batch { shots: 0, hist: Vec::new(), s1: vec![0; modes], s2: vec![0; modes * modes], max_count: vec![0; modes], negative: 0, patterns: HashMap::new() };
    let mut nz: Vec<(usize, u64)> = Vec::with_capacity(modes);
    for row in rows.chunks(modes) {
        b.shots += 1;
        nz.clear();
        let mut total = 0usize;
        let mut clean = true;
        for (i, &v) in row.iter().enumerate() {
            if v > 127 {
                // A negative entry: counted, and the sample kept out of the pattern table.
                b.negative += 1;
                clean = false;
                continue;
            }
            if v > 0 {
                nz.push((i, u64::from(v)));
                total += v as usize;
                b.max_count[i] = b.max_count[i].max(v);
            }
        }
        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 &(j, nj) in &nz[x..] {
                b.s2[i * modes + j] += ni * nj;
            }
        }
        if let Some(n) = max_n
            && total <= n
            && clean
        {
            *b.patterns.entry(common::pack(row)).or_insert(0) += 1;
        }
    }
    b
}

/// The mean and standard error of a statistic across batches.
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 the observed pair cumulants compare with a hypothesis's: the root mean
/// square and the share within 3 of the standardised differences, the
/// correlation, and the least-squares slope of observed on predicted.
fn cumulant_agreement(obs: &[f64], se: &[f64], hyp: &[f64]) -> Value {
    let k = obs.len() as f64;
    let z: Vec<f64> = obs.iter().zip(se).zip(hyp).map(|((o, s), h)| (o - h) / s).collect();
    let rms = (z.iter().map(|x| x * x).sum::<f64>() / k).sqrt();
    let within3 = z.iter().filter(|x| x.abs() <= 3.0).count() as f64 / k;
    let (mo, mh) = (obs.iter().sum::<f64>() / k, hyp.iter().sum::<f64>() / k);
    let (mut sxy, mut sxx, mut syy, mut soh, mut shh) = (0.0, 0.0, 0.0, 0.0, 0.0);
    for (o, h) in obs.iter().zip(hyp) {
        sxy += (o - mo) * (h - mh);
        sxx += (o - mo) * (o - mo);
        syy += (h - mh) * (h - mh);
        soh += o * h;
        shh += h * h;
    }
    let corr = if syy > 0.0 { Some(sxy / (sxx * syy).sqrt()) } else { None };
    let slope = if shh > 0.0 { Some(soh / shh) } else { None };
    json!({ "rms_z": rms, "share_within_3": within3, "correlation": corr, "slope": slope, "mean_abs_diff": obs.iter().zip(hyp).map(|(o, h)| (o - h).abs()).sum::<f64>() / k })
}

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
}

/// The coherent hypothesis's total: Poisson of mean `mu`, normalised over
/// `0..=nmax` (with `nmax` far past the mean, what is cut is below rounding).
fn poisson(mu: f64, nmax: usize) -> Vec<f64> {
    let mut p = vec![1.0];
    for n in 1..=nmax {
        let last = p[n - 1];
        p.push(last * mu / n as f64);
    }
    let s: f64 = p.iter().sum();
    p.iter().map(|x| x / s).collect()
}

/// Convolution of distributions over photon numbers, cut at `nmax`.
fn convolve(a: &[f64], b: &[f64], nmax: usize) -> Vec<f64> {
    let mut out = vec![0.0; nmax + 1];
    for (i, x) in a.iter().enumerate() {
        for (j, y) in b.iter().enumerate() {
            if i + j <= nmax {
                out[i + j] += x * y;
            }
        }
    }
    out
}

#[allow(clippy::too_many_lines)]
fn run(args: &Args, modes: usize) -> (Value, Value) {
    let set = args.set.as_str();
    let path = |f: &str| format!("{set}/{f}");
    let root: &Path = &args.data;
    let t0 = Instant::now();
    let params = common::read_pinned(root, &path("program_params.json"));
    let t_pub = common::complex_matrix(&common::read_pinned(root, &path("T.npy")));
    let u_pub = common::complex_matrix(&common::read_pinned(root, &path("U.npy")));
    let r_npy = common::real_vector(&common::read_pinned(root, &path("r.npy")));
    let sample_bytes = common::read_pinned(root, &path("samples.npy"));
    let read_secs = t0.elapsed().as_secs_f64();
    assert_eq!(t_pub.n, modes, "the transfer matrix has the set's modes");
    let m = modes;
    println!("{set}: {m} modes; files read and checked against their pins in {read_secs:.1} s");

    // 1. The transfer matrix and unitary, rebuilt from the program.
    let prog = common::program(&params);
    let (dt, qt) = common::distance_up_to_quarter_phase(&prog.transfer(m, false), &t_pub);
    let (du, qu) = common::distance_up_to_quarter_phase(&prog.transfer(m, true), &u_pub);
    println!("rebuilt from the loop program: max |T − published| {dt:.2e}, max |U − published| {du:.2e} (global phase i^{qt}, i^{qu})");
    // The squeezing the published ground truth pairs with T's columns, and where it sits in
    // the program's list (the two agree to the last bit or two, not bit for bit).
    let r_gt: Vec<f64> = r_npy[..m].to_vec();
    let offset = (0..=prog.r.len() - m).find(|&o| (0..m).all(|i| (prog.r[o + i] - r_gt[i]).abs() < 1e-12));

    // 2. The hypotheses.
    let nb: Vec<f64> = r_gt.iter().map(|&r| nbar(r)).collect();
    let gt = GaussianState::through(&r_gt.iter().map(|&r| Input::Squeezed(r)).collect::<Vec<_>>(), &t_pub).expect("the ground truth");
    let thermal = GaussianState::through(&nb.iter().map(|&n| Input::Thermal(n)).collect::<Vec<_>>(), &t_pub).expect("thermal");
    let squashed = GaussianState::through(&nb.iter().map(|&n| Input::Squashed(n)).collect::<Vec<_>>(), &t_pub).expect("squashed");
    // The ground truth with the program's first m squeezing values, where they differ from r.npy's.
    let alt = (offset != Some(0)).then(|| {
        let r0: Vec<f64> = prog.r[..m].to_vec();
        (r0.clone(), GaussianState::through(&r0.iter().map(|&r| Input::Squeezed(r)).collect::<Vec<_>>(), &t_pub).expect("the alternative"))
    });

    // 3. The samples, counted in batches.
    let t1 = Instant::now();
    let (shots, sm, rows) = common::sample_rows(&sample_bytes);
    assert_eq!(sm, m, "the samples have the set's modes");
    let per = shots.div_ceil(args.batches);
    let pattern_cap = (m <= 16).then_some(args.max_n);
    let batches: Vec<Batch> = parallel(args.batches, args.threads, |b| {
        let lo = (b * per).min(shots);
        let hi = ((b + 1) * per).min(shots);
        count_batch(&rows[lo * m..hi * m], m, pattern_cap)
    });
    let count_secs = t1.elapsed().as_secs_f64();
    let k = batches.iter().map(|b| b.shots).sum::<u64>();
    assert_eq!(k as usize, shots);
    let negative: u64 = batches.iter().map(|b| b.negative).sum();
    let nmax_obs = batches.iter().map(|b| b.hist.len()).max().unwrap() - 1;
    let mut hist = vec![0u64; nmax_obs + 1];
    let mut s1 = vec![0u64; m];
    let mut s2 = vec![0u64; m * m];
    let mut max_count = vec![0u8; m];
    for b in &batches {
        for (i, x) in b.hist.iter().enumerate() {
            hist[i] += x;
        }
        for i in 0..m {
            s1[i] += b.s1[i];
            max_count[i] = max_count[i].max(b.max_count[i]);
        }
        for (x, y) in s2.iter_mut().zip(&b.s2) {
            *x += y;
        }
    }
    let kf = k as f64;
    println!("{k} samples counted in {count_secs:.1} s; most photons in one sample {nmax_obs}; most in one detector {}; negative entries {negative}", max_count.iter().max().unwrap());

    // 4a. The total photon number.
    // Far enough past the largest observed total that the exact tail beyond is below rounding.
    let nmax = nmax_obs + nmax_obs / 2 + 60;
    let t2 = Instant::now();
    let p_gt = gt.photon_number_distribution(nmax).expect("the ground truth's total");
    let p_th = thermal.photon_number_distribution(nmax).expect("thermal total");
    let p_sq = squashed.photon_number_distribution(nmax).expect("squashed total");
    let (mean_gt, cov_gt) = gt.photon_moments();
    let (_, cov_th) = thermal.photon_moments();
    let (_, cov_sq) = squashed.photon_moments();
    let mu_gt: f64 = mean_gt.iter().sum();
    let var_gt: f64 = cov_gt.iter().sum();
    let p_coh = poisson(mu_gt, nmax);
    // Distinguishable squeezed light: each input on its own, its total that of
    // one squeezed mode through its column's transmittance.
    let mut p_di = vec![1.0];
    let mut cov_di = vec![0.0; m * m];
    for col in 0..m {
        let eta: f64 = (0..m).map(|i| t_pub.at(i, col).re * t_pub.at(i, col).re + t_pub.at(i, col).im * t_pub.at(i, col).im).sum();
        let one = GaussianState::through(&[Input::Squeezed(r_gt[col])], &Mat { n: 1, a: vec![wai_quantum::quantum_gbs::C { re: eta.sqrt(), im: 0.0 }] }).expect("one input");
        p_di = convolve(&p_di, &one.photon_number_distribution(nmax).expect("one input's total"), nmax);
        // Its pair cumulants: |T_ik T_jk|² (s⁴ + s² c²), with s² = n̄ and c² = n̄ + 1.
        let s2k = nb[col];
        for i in 0..m {
            for j in 0..m {
                let (a, b) = (t_pub.at(i, col), t_pub.at(j, col));
                let w = (a.re * a.re + a.im * a.im) * (b.re * b.re + b.im * b.im);
                cov_di[i * m + j] += if i == j { 0.0 } else { w * (s2k * s2k + s2k * (s2k + 1.0)) };
            }
        }
    }
    let exact_secs = t2.elapsed().as_secs_f64();
    let h_obs: Vec<f64> = hist.iter().map(|&x| x as f64 / kf).collect();
    let sum1: f64 = p_gt.iter().sum();
    let pm1: f64 = p_gt.iter().enumerate().map(|(n, x)| n as f64 * x).sum();
    let pm2: f64 = p_gt.iter().enumerate().map(|(n, x)| (n * n) as f64 * x).sum();
    let obs_mean = hist.iter().enumerate().map(|(n, &c)| n as f64 * c as f64).sum::<f64>() / kf;
    let obs_var = hist.iter().enumerate().map(|(n, &c)| (n as f64 - obs_mean) * (n as f64 - obs_mean) * c as f64).sum::<f64>() / (kf - 1.0);
    let mode_obs = (0..hist.len()).max_by_key(|&n| (hist[n], std::cmp::Reverse(n))).unwrap();
    // What K draws from the exact total alone give.
    let floor: Vec<f64> = parallel(args.replicas, args.threads, |rep| {
        let mut rng = common::Rng(0x5eed_0000 + rep as u64);
        let c = common::draw(&p_gt, k, &mut rng);
        tvd(&c.iter().map(|&x| x as f64 / kf).collect::<Vec<_>>(), &p_gt)
    });
    let (floor_mean, floor_se) = mean_se(&floor);
    let floor_sd = floor_se * (args.replicas as f64).sqrt();
    let totals = json!({
        "observed": { "samples": k, "histogram": hist, "mean": obs_mean, "mean_se": (obs_var / kf).sqrt(), "variance": obs_var, "most_frequent": mode_obs, "max": nmax_obs },
        "ground_truth": { "mean": mu_gt, "variance": var_gt, "most_probable": (0..=nmax).max_by(|&a, &b| p_gt[a].total_cmp(&p_gt[b])).unwrap(),
            "distribution": p_gt, "check": { "sum": sum1, "mean_from_distribution": pm1, "variance_from_distribution": pm2 - pm1 * pm1 } },
        "tvd": { "ground_truth": tvd(&h_obs, &p_gt), "thermal": tvd(&h_obs, &p_th), "squashed": tvd(&h_obs, &p_sq), "coherent": tvd(&h_obs, &p_coh), "distinguishable": tvd(&h_obs, &p_di) },
        "tvd_from_sampling_alone": { "replicas": args.replicas, "mean": floor_mean, "sd": floor_sd },
    });
    println!("total photons: observed mean {obs_mean:.4} ± {:.4}, exact {mu_gt:.4}; TVD to exact {:.4} (sampling alone {floor_mean:.4} ± {floor_sd:.4}); thermal {:.4}, squashed {:.4}, coherent {:.4}, distinguishable {:.4}",
        (obs_var / kf).sqrt(), totals["tvd"]["ground_truth"].as_f64().unwrap(), totals["tvd"]["thermal"].as_f64().unwrap(), totals["tvd"]["squashed"].as_f64().unwrap(), totals["tvd"]["coherent"].as_f64().unwrap(), totals["tvd"]["distinguishable"].as_f64().unwrap());

    // 4b. First-order cumulants: every mode's mean.
    let obs_means: Vec<f64> = s1.iter().map(|&x| x as f64 / kf).collect();
    let z1: Vec<f64> = (0..m).map(|i| (obs_means[i] - mean_gt[i]) / (cov_gt[i * m + i] / kf).sqrt()).collect();
    let rel1: Vec<f64> = (0..m).map(|i| obs_means[i] / mean_gt[i] - 1.0).collect();
    let firsts = json!({
        "observed": obs_means, "ground_truth": mean_gt,
        "max_abs_z": z1.iter().fold(0.0f64, |a, z| a.max(z.abs())),
        "rms_z": (z1.iter().map(|z| z * z).sum::<f64>() / m as f64).sqrt(),
        "mean_relative_deviation": rel1.iter().sum::<f64>() / m as f64,
        "max_abs_relative_deviation": rel1.iter().fold(0.0f64, |a, x| a.max(x.abs())),
    });
    println!("mode means: rms z {:.1}, max |z| {:.1}, mean relative deviation {:+.4}, largest {:.4}", firsts["rms_z"].as_f64().unwrap(), firsts["max_abs_z"].as_f64().unwrap(), firsts["mean_relative_deviation"].as_f64().unwrap(), firsts["max_abs_relative_deviation"].as_f64().unwrap());

    // 4c. Second-order cumulants of every pair, with standard errors from the batches.
    let pairs: Vec<(usize, usize)> = (0..m).flat_map(|i| (i + 1..m).map(move |j| (i, j))).collect();
    let obs_c: Vec<f64> = pairs.iter().map(|&(i, j)| s2[i * m + j] as f64 / kf - (s1[i] as f64 / kf) * (s1[j] as f64 / kf)).collect();
    let se_c: Vec<f64> = pairs
        .iter()
        .map(|&(i, j)| {
            let per_batch: Vec<f64> = batches.iter().map(|b| {
                let n = b.shots as f64;
                b.s2[i * m + j] as f64 / n - (b.s1[i] as f64 / n) * (b.s1[j] as f64 / n)
            }).collect();
            mean_se(&per_batch).1
        })
        .collect();
    let hyp = |c: &[f64]| -> Vec<f64> { pairs.iter().map(|&(i, j)| c[i * m + j]).collect() };
    let zeros = vec![0.0; pairs.len()];
    let median_se = {
        let mut s = se_c.clone();
        s.sort_by(f64::total_cmp);
        s[s.len() / 2]
    };
    let seconds = json!({
        "pairs": pairs.len(), "batches": args.batches,
        "ground_truth": cumulant_agreement(&obs_c, &se_c, &hyp(&cov_gt)),
        "thermal": cumulant_agreement(&obs_c, &se_c, &hyp(&cov_th)),
        "squashed": cumulant_agreement(&obs_c, &se_c, &hyp(&cov_sq)),
        "coherent": cumulant_agreement(&obs_c, &se_c, &zeros),
        "distinguishable": cumulant_agreement(&obs_c, &se_c, &hyp(&cov_di)),
        "median_se": median_se,
    });
    for h in ["ground_truth", "thermal", "squashed", "coherent", "distinguishable"] {
        let x = &seconds[h];
        println!("pair cumulants vs {h:15}: rms z {:8.2}, within 3σ {:.3}, correlation {}, slope {}", x["rms_z"].as_f64().unwrap(), x["share_within_3"].as_f64().unwrap(), x["correlation"].as_f64().map_or("—".into(), |v| format!("{v:.4}")), x["slope"].as_f64().map_or("—".into(), |v| format!("{v:.4}")));
    }

    let mut result = json!({
        "set": set, "modes": m,
        "rebuilt_from_program": { "max_abs_diff_T": dt, "max_abs_diff_U": du, "global_phase_quarter_turns_T": qt, "global_phase_quarter_turns_U": qu },
        "squeezing": { "r_npy_length": r_npy.len(), "offset_in_program_list": offset, "program_list_length": prog.r.len() },
        "totals": totals, "first_order": firsts, "second_order": seconds,
        "phases_seconds": { "read_and_hash": read_secs, "count_samples": count_secs, "exact_totals_and_cumulants": exact_secs },
    });
    if let Some((r0, st)) = &alt {
        let (mean_alt, _) = st.photon_moments();
        let z: Vec<f64> = (0..m).map(|i| (obs_means[i] - mean_alt[i]) / (cov_gt[i * m + i] / kf).sqrt()).collect();
        result["alternative_squeezing"] = json!({
            "why": "the program's first m squeezing values, which the other four sets pair with their inputs, where r.npy holds others",
            "r": r0,
            "first_order_rms_z": (z.iter().map(|x| x * x).sum::<f64>() / m as f64).sqrt(),
        });
    }

    // 4d. The 16-mode set: every pattern of up to max_n photons.
    if m <= 16 {
        let t3 = Instant::now();
        let mut merged: HashMap<u64, u64> = HashMap::new();
        for b in &batches {
            for (key, c) in &b.patterns {
                *merged.entry(*key).or_insert(0) += c;
            }
        }
        let tables: Vec<(&str, PhotonTable)> = {
            let mut v = vec![("ground_truth", gt.photon_table().unwrap()), ("thermal", thermal.photon_table().unwrap()), ("squashed", squashed.photon_table().unwrap())];
            if let Some((_, st)) = &alt {
                v.push(("alternative_squeezing", st.photon_table().unwrap()));
            }
            v
        };
        let mut sectors = Vec::new();
        for n in 0..=args.max_n {
            let pats = common::patterns(m, n);
            let observed: Vec<u64> = pats.iter().map(|p| merged.get(&common::pack(p)).copied().unwrap_or(0)).collect();
            let kn: u64 = observed.iter().sum();
            let q: Vec<f64> = observed.iter().map(|&c| c as f64 / kn.max(1) as f64).collect();
            let mut row = json!({ "n": n, "patterns": pats.len(), "samples": kn, "patterns_seen": observed.iter().filter(|&&c| c > 0).count() });
            let mut p_ground = Vec::new();
            for (name, table) in &tables {
                let p: Vec<f64> = parallel(pats.len(), args.threads, |i| table.probability(&pats[i].iter().map(|&x| x as usize).collect::<Vec<_>>()).expect("a pattern"));
                let s: f64 = p.iter().sum();
                let pn: Vec<f64> = p.iter().map(|x| x / s).collect();
                let (f, t) = common::fidelity_tvd(&pn, &q);
                row[*name] = json!({ "fidelity": f, "tvd": t, "sector_probability": s, "fidelity_bits": common::bits(f), "tvd_bits": common::bits(t) });
                if *name == "ground_truth" {
                    row["exact_sector_probability"] = json!(p_gt[n]);
                    p_ground = pn;
                }
            }
            // Coherent light with the same mode means: a multinomial within the sector.
            let mut pc: Vec<f64> = pats.iter().map(|p| p.iter().enumerate().fold(1.0, |acc, (i, &c)| {
                let mut v = acc;
                for d in 1..=c as usize {
                    v *= mean_gt[i] / d as f64;
                }
                v
            })).collect();
            let s: f64 = pc.iter().sum();
            pc.iter_mut().for_each(|x| *x /= s);
            let (f, t) = common::fidelity_tvd(&pc, &q);
            row["coherent"] = json!({ "fidelity": f, "tvd": t });
            // What kn draws from the exact distribution alone give.
            let reps: Vec<(f64, f64)> = parallel(args.replicas, args.threads, |rep| {
                let mut rng = common::Rng(0x0b5e_0000 + 1000 * n as u64 + rep as u64);
                let c = common::draw(&p_ground, kn, &mut rng);
                common::fidelity_tvd(&p_ground, &c.iter().map(|&x| x as f64 / kn.max(1) as f64).collect::<Vec<_>>())
            });
            let (fm, fse) = mean_se(&reps.iter().map(|x| x.0).collect::<Vec<_>>());
            let (tm, tse) = mean_se(&reps.iter().map(|x| x.1).collect::<Vec<_>>());
            let sd = (args.replicas as f64).sqrt();
            row["from_sampling_alone"] = json!({ "replicas": args.replicas, "fidelity_mean": fm, "fidelity_sd": fse * sd, "tvd_mean": tm, "tvd_sd": tse * sd });
            println!(
                "N = {n}: {kn:>9} samples over {:>6} patterns · ground truth F {:.4} TVD {:.4} (sampling alone F {fm:.4} TVD {tm:.4} ± {:.4}) · thermal TVD {:.4} · squashed {:.4} · coherent {:.4}{}",
                pats.len(),
                row["ground_truth"]["fidelity"].as_f64().unwrap(),
                row["ground_truth"]["tvd"].as_f64().unwrap(),
                tse * sd,
                row["thermal"]["tvd"].as_f64().unwrap(),
                row["squashed"]["tvd"].as_f64().unwrap(),
                row["coherent"]["tvd"].as_f64().unwrap(),
                row.get("alternative_squeezing").map_or(String::new(), |a| format!(" · alternative squeezing TVD {:.4}", a["tvd"].as_f64().unwrap()))
            );
            sectors.push(row);
        }
        result["patterns"] = json!({ "sectors": sectors, "seconds": t3.elapsed().as_secs_f64() });
        if let Some(path) = &args.counts {
            let text = common::counts_text(m, 4.min(args.max_n), |p| merged.get(&common::pack(p)).copied().unwrap_or(0));
            std::fs::write(path, &text).unwrap_or_else(|e| panic!("writing {}: {e}", path.display()));
            result["pattern_counts_file"] = json!({ "max_n": 4.min(args.max_n), "sha256": common::hex(&common::sha256(text.as_bytes())) });
            println!("pattern counts for N ≤ 4 written to {}", path.display());
        }
    }

    // 4e. The paper's per-N scores, over the first K samples with N photons.
    if let Some((lo, hi)) = args.scores {
        let t4 = Instant::now();
        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() < args.per_n {
                by_n[n].push(row.iter().map(|&v| v as usize).collect());
            }
        }
        let t_gt = gt.photon_table().unwrap();
        let t_th = thermal.photon_table().unwrap();
        let t_sq = squashed.photon_table().unwrap();
        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 rows_out = Vec::new();
        for n in lo..=hi {
            let samples = &by_n[n];
            if samples.is_empty() {
                continue;
            }
            // From N = 12 each hafnian runs on every core itself.
            let outer = if n >= 12 { 1 } else { args.threads };
            let probs: Vec<(f64, f64, f64)> = parallel(samples.len(), outer, |i| {
                (t_gt.probability(&samples[i]).unwrap(), t_th.probability(&samples[i]).unwrap(), t_sq.probability(&samples[i]).unwrap())
            });
            // Cross entropy: ln(Pr(S) / (Pr(N) / C(N+M−1, N))).
            let ln_uniform = p_gt[n].ln() - ln_binom(n + m - 1, n);
            let xe: Vec<f64> = probs.iter().map(|p| p.0.ln() - ln_uniform).collect();
            // Bayesian log-average against each hypothesis: ln(Pr₀(S) Pr_I(N) / (Pr_I(S) Pr₀(N))).
            let coh = |s: &[usize]| -> f64 {
                // Multinomial ln: ln N! − Σ ln nᵢ! + Σ nᵢ ln(μᵢ/μ).
                let mut v: f64 = (1..=n).map(|d| (d as f64).ln()).sum();
                for (i, &c) in s.iter().enumerate() {
                    v -= (1..=c).map(|d| (d as f64).ln()).sum::<f64>();
                    v += c as f64 * (mean_gt[i] / mu_gt).ln();
                }
                v
            };
            let b_th: Vec<f64> = probs.iter().map(|p| (p.0 / p_gt[n]).ln() - (p.1 / p_th[n]).ln()).collect();
            let b_sq: Vec<f64> = probs.iter().map(|p| (p.0 / p_gt[n]).ln() - (p.2 / p_sq[n]).ln()).collect();
            let b_co: Vec<f64> = probs.iter().zip(samples).map(|(p, s)| (p.0 / p_gt[n]).ln() - coh(s)).collect();
            let ms = |v: &[f64]| -> Value {
                let (mean, se) = if v.len() > 1 { let k = v.len() as f64; let mu = v.iter().sum::<f64>() / k; (mu, (v.iter().map(|x| (x - mu) * (x - mu)).sum::<f64>() / (k - 1.0) / k).sqrt()) } else { (v[0], f64::NAN) };
                json!({ "mean": mean, "se": se })
            };
            let row = json!({ "n": n, "samples": samples.len(), "cross_entropy": ms(&xe), "bayes_vs_thermal": ms(&b_th), "bayes_vs_squashed": ms(&b_sq), "bayes_vs_coherent": ms(&b_co) });
            println!(
                "N = {n}: {} samples · cross entropy {:.4} ± {:.4} · Bayes vs thermal {:.3} ± {:.3}, squashed {:.3} ± {:.3}, coherent {:.3} ± {:.3}",
                samples.len(),
                row["cross_entropy"]["mean"].as_f64().unwrap(), row["cross_entropy"]["se"].as_f64().unwrap_or(f64::NAN),
                row["bayes_vs_thermal"]["mean"].as_f64().unwrap(), row["bayes_vs_thermal"]["se"].as_f64().unwrap_or(f64::NAN),
                row["bayes_vs_squashed"]["mean"].as_f64().unwrap(), row["bayes_vs_squashed"]["se"].as_f64().unwrap_or(f64::NAN),
                row["bayes_vs_coherent"]["mean"].as_f64().unwrap(), row["bayes_vs_coherent"]["se"].as_f64().unwrap_or(f64::NAN),
            );
            rows_out.push(row);
        }
        result["scores"] = json!({ "per_n_cap": args.per_n, "selection": "the first samples in file order with each N", "rows": rows_out, "seconds": t4.elapsed().as_secs_f64(),
            "logarithms": "the platform's f64 ln: these scores are not bit-identical across platforms, unlike the rest of the result" });
    }

    let mut sources: Vec<Value> = ["program_params.json", "T.npy", "U.npy", "r.npy", "samples.npy"].iter().map(|f| common::source(&path(f))).collect();
    sources.push(json!({ "repository": common::REPOSITORY, "commit": common::COMMIT, "path": "LICENSE", "sha256": common::LICENSE_SHA256, "licence": "Apache-2.0" }));
    let request = json!({
        "experiment": "time-domain photonic Gaussian boson sampling: Madsen et al., Nature 606, 75 (2022), doi:10.1038/s41586-022-04725-x",
        "set": set,
        "source": sources,
        "ground_truth": "squeezed vacua with r.npy's first m values through the published T (the paper's (r, T))",
        "hypotheses": {
            "thermal": "thermal states with the same mean photon number per input, through T",
            "squashed": "squashed states (vacuum noise in one quadrature) with the same mean photon number per input, through T",
            "coherent": "coherent light with the ground truth's mean in every mode: independent Poisson counts",
            "distinguishable": "each squeezed input on its own through T, the outputs added",
        },
        "method": {
            "transfer": "quantum_gbs::Loops from program_params.json (delays 1, 6, 36; detection common × channel (bin mod 16))",
            "patterns": "quantum_gbs::PhotonTable: Haf(A_n)/(Π nᵢ! √det Q) for every pattern of up to max_n photons",
            "totals": "quantum_gbs::GaussianState::photon_number_distribution",
            "cumulants": "quantum_gbs::GaussianState::photon_moments",
            "standard_errors": "spread of the statistic across contiguous batches of samples",
            "sampling_alone": "draws of the observed sample count from the exact distribution, SplitMix64 seeds 0x0b5e0000 + 1000 N + replica (patterns) and 0x5eed0000 + replica (totals)",
        },
        "settings": { "max_n": args.max_n, "replicas": args.replicas, "batches": args.batches, "scores": args.scores.map(|(a, b)| json!([a, b])), "per_n": args.per_n },
        "published": published(set),
    });
    (request, result)
}

/// What the paper states for each set, as its text gives it.
fn published(set: &str) -> Value {
    match set {
        "fig2" => json!({ "samples": "84.1 × 10⁶", "photon_numbers_compared": "3, 4, 5 and 6", "claim": "fidelities in excess of 99% and TVDs below or equal to 6.5% for the total photon-number sectors considered", "normalisation": "pᵢ and qᵢ normalised by the probability of their total photon number" }),
        "fig3a" => json!({ "mean_photons": "21.120 ± 0.006", "modes": 216, "scores": "cross entropy over 10,000 samples per N, N = 10…26, above every classical spoofer" }),
        "fig3b" => json!({ "mean_photons": "22.416 ± 0.006", "modes": 72, "scores": "Bayesian log-average over 2,000 samples per N, N = 10…26, above zero against squashed, thermal, coherent and distinguishable squeezed hypotheses" }),
        "fig4" => json!({ "samples": "10⁶", "mean_photons": "centred at 125 (every hypothesis has the ground truth's mode means)", "most_probable_observed": "N = 124.35 ± 0.02", "max_observed": 219, "claim": "the ground truth is the only hypothesis that explains the pair cumulants well" }),
        "figS15" => json!({ "modes": 288, "note": "shown in the Supplementary Information; no value is stated in the main text" }),
        _ => Value::Null,
    }
}