wai-quantum 0.3.30

//! Gaussian boson sampling, exactly — `wai.quantum.gbs`.
//!
//! A Gaussian boson sampler sends squeezed light through an interferometer
//! and counts photons at the outputs. The probability of each outcome is a
//! matrix function of the Gaussian state:
//! - a **hafnian** for photon-number-resolving detectors;
//! - a **torontonian** for threshold (click/no-click) detectors.
//!
//! Both are #P-hard in general. That hardness is the basis of every photonic
//! quantum-advantage claim, and computing them exactly is how a claim is
//! checked wherever it can be.
//!
//! This module computes them exactly and reproducibly.
//!
//! - [`hafnian`] uses power traces: a sum over `2^{n/2}` subsets of pairs, each
//!   the `λ^{n/2}` coefficient of `exp(Σ tr(Mᵏ) λᵏ / 2k)`. The traces come from
//!   a Hessenberg reduction and La Budde's characteristic polynomial, so the
//!   total cost is `O(n³ 2^{n/2})`.
//! - [`torontonian`] is the inclusion–exclusion sum over `2^m` subsets of the
//!   clicked modes, each a determinant.
//! - [`GaussianState`] is a state built from single-mode squeezers, an
//!   interferometer and uniform loss. It is held as its complex covariance
//!   matrix in the `(a, a†)` basis, with vacuum `I/2`.
//! - [`GaussianState::click_probability`] and
//!   [`GaussianState::photon_probability`] give exact pattern probabilities.
//! - [`GaussianState::sample_clicks`] draws exact samples mode by mode from
//!   marginal torontonians.
//! - [`GaussianState::click_marginals`] gives the one- and two-mode click
//!   statistics that any sampler, a device included, must reproduce. They can
//!   be computed for any number of modes.
//!
//! Every routine uses `+ − × ÷ √` only, in a fixed order. Hafnian and
//! torontonian sums are split into a fixed number of chunks, threaded natively
//! and run in sequence on targets without threads, and added in chunk order.
//! Probabilities and samples are therefore the same bits on every machine.
//!
//! # Checked against an independent implementation
//!
//! - **Click probabilities:** every pattern of three lossy states of four to six
//!   modes agrees to 1.1·10⁻¹⁰, the worst being the smallest probabilities.
//! - **Photon-number probabilities:** every pattern of up to four photons
//!   agrees to 2.3·10⁻¹³.
//! - **Hafnians** of random complex symmetric matrices agree to 1.5·10⁻¹² up to
//!   `20 × 20` and to 4.5·10⁻⁸ at `36 × 36`, where the alternating
//!   `2^{n/2}`-term sum loses precision in either implementation. A `36 × 36`
//!   hafnian takes under two seconds on one machine.
//!
//! # What it checks
//!
//! [`GaussianState::check_samples`] holds observed click patterns, from a
//! device or from any sampler, to the state's exact one- and two-mode click
//! statistics. Exact samples of the state pass. The same samples tested against
//! a different interferometer fail by a wide margin.
//!
//! These statistics can be computed for any number of modes. Exact pattern
//! probabilities cost `2^{clicks}` determinants, so they are practical up to
//! a few dozen clicks.

pub use crate::linalg::C;

fn add(a: C, b: C) -> C {
    C { re: a.re + b.re, im: a.im + b.im }
}
fn sub(a: C, b: C) -> C {
    C { re: a.re - b.re, im: a.im - b.im }
}
fn mul(a: C, b: C) -> C {
    C { re: a.re * b.re - a.im * b.im, im: a.re * b.im + a.im * b.re }
}
fn div(a: C, b: C) -> C {
    let d = b.re * b.re + b.im * b.im;
    C { re: (a.re * b.re + a.im * b.im) / d, im: (a.im * b.re - a.re * b.im) / d }
}
fn scale(a: C, s: f64) -> C {
    C { re: a.re * s, im: a.im * s }
}
fn conj(a: C) -> C {
    C { re: a.re, im: -a.im }
}
fn abs2(a: C) -> f64 {
    a.re * a.re + a.im * a.im
}
const ZERO: C = C { re: 0.0, im: 0.0 };
const ONE: C = C { re: 1.0, im: 0.0 };

/// A dense complex matrix, row-major.
#[derive(Clone, Debug)]
pub struct Mat {
    pub n: usize,
    pub a: Vec<C>,
}

impl Mat {
    pub fn zeros(n: usize) -> Mat {
        Mat { n, a: vec![ZERO; n * n] }
    }
    pub fn identity(n: usize) -> Mat {
        let mut m = Mat::zeros(n);
        for i in 0..n {
            m.a[i * n + i] = ONE;
        }
        m
    }
    #[inline]
    pub fn at(&self, i: usize, j: usize) -> C {
        self.a[i * self.n + j]
    }
    #[inline]
    fn set(&mut self, i: usize, j: usize, v: C) {
        self.a[i * self.n + j] = v;
    }
    fn mul(&self, o: &Mat) -> Mat {
        let n = self.n;
        let mut out = Mat::zeros(n);
        for i in 0..n {
            for k in 0..n {
                let x = self.at(i, k);
                if x.re == 0.0 && x.im == 0.0 {
                    continue;
                }
                for j in 0..n {
                    out.a[i * n + j] = add(out.a[i * n + j], mul(x, o.at(k, j)));
                }
            }
        }
        out
    }
    fn dagger(&self) -> Mat {
        let n = self.n;
        let mut out = Mat::zeros(n);
        for i in 0..n {
            for j in 0..n {
                out.set(j, i, conj(self.at(i, j)));
            }
        }
        out
    }
    /// The principal submatrix on `idx` (repeats allowed).
    pub fn sub(&self, idx: &[usize]) -> Mat {
        let m = idx.len();
        let mut out = Mat::zeros(m);
        for (r, &i) in idx.iter().enumerate() {
            for (c, &j) in idx.iter().enumerate() {
                out.a[r * m + c] = self.at(i, j);
            }
        }
        out
    }

    /// Determinant by LU decomposition with partial pivoting.
    pub fn det(&self) -> C {
        let n = self.n;
        let mut m = self.a.clone();
        let mut d = ONE;
        for col in 0..n {
            let mut piv = col;
            for r in col + 1..n {
                if abs2(m[r * n + col]) > abs2(m[piv * n + col]) {
                    piv = r;
                }
            }
            if abs2(m[piv * n + col]) == 0.0 {
                return ZERO;
            }
            if piv != col {
                for c in 0..n {
                    m.swap(piv * n + c, col * n + c);
                }
                d = scale(d, -1.0);
            }
            let p = m[col * n + col];
            d = mul(d, p);
            for r in col + 1..n {
                let f = div(m[r * n + col], p);
                if f.re == 0.0 && f.im == 0.0 {
                    continue;
                }
                for c in col..n {
                    m[r * n + c] = sub(m[r * n + c], mul(f, m[col * n + c]));
                }
            }
        }
        d
    }

    /// Inverse by Gauss–Jordan elimination with partial pivoting.
    pub fn inverse(&self) -> Option<Mat> {
        let n = self.n;
        let mut m = self.a.clone();
        let mut inv = Mat::identity(n).a;
        for col in 0..n {
            let mut piv = col;
            for r in col + 1..n {
                if abs2(m[r * n + col]) > abs2(m[piv * n + col]) {
                    piv = r;
                }
            }
            if abs2(m[piv * n + col]) == 0.0 {
                return None;
            }
            for c in 0..n {
                m.swap(piv * n + c, col * n + c);
                inv.swap(piv * n + c, col * n + c);
            }
            let p = m[col * n + col];
            for c in 0..n {
                m[col * n + c] = div(m[col * n + c], p);
                inv[col * n + c] = div(inv[col * n + c], p);
            }
            for r in 0..n {
                if r == col {
                    continue;
                }
                let f = m[r * n + col];
                if f.re == 0.0 && f.im == 0.0 {
                    continue;
                }
                for c in 0..n {
                    m[r * n + c] = sub(m[r * n + c], mul(f, m[col * n + c]));
                    inv[r * n + c] = sub(inv[r * n + c], mul(f, inv[col * n + c]));
                }
            }
        }
        Some(Mat { n, a: inv })
    }
}

/// `tr(Mᵏ)` for `k = 1..=kmax`: reduce to Hessenberg form by Householder
/// reflections, take the characteristic polynomial by La Budde's recurrence,
/// and turn its coefficients into power sums by Newton's identities.
fn power_traces(m: &Mat, kmax: usize) -> Vec<C> {
    let n = m.n;
    if n == 0 {
        return vec![ZERO; kmax];
    }
    let mut h = m.a.clone();
    // Householder reduction to upper Hessenberg form (similarity).
    for col in 0..n.saturating_sub(2) {
        let mut norm2 = 0.0;
        for r in col + 1..n {
            norm2 += abs2(h[r * n + col]);
        }
        let x0 = h[(col + 1) * n + col];
        let alpha = norm2.sqrt();
        if alpha == 0.0 {
            continue;
        }
        let ax0 = abs2(x0).sqrt();
        // v = x + e^{i arg x0} alpha e_1, normalised.
        let phase = if ax0 == 0.0 { ONE } else { scale(x0, 1.0 / ax0) };
        let mut v: Vec<C> = (col + 1..n).map(|r| h[r * n + col]).collect();
        v[0] = add(v[0], scale(phase, alpha));
        let vn = v.iter().map(|x| abs2(*x)).sum::<f64>().sqrt();
        if vn == 0.0 {
            continue;
        }
        for x in v.iter_mut() {
            *x = scale(*x, 1.0 / vn);
        }
        // H ← (I − 2vv†) H (I − 2vv†), on rows/cols col+1..n.
        for c in 0..n {
            let mut s = ZERO;
            for (k, vk) in v.iter().enumerate() {
                s = add(s, mul(conj(*vk), h[(col + 1 + k) * n + c]));
            }
            for (k, vk) in v.iter().enumerate() {
                let idx = (col + 1 + k) * n + c;
                h[idx] = sub(h[idx], scale(mul(*vk, s), 2.0));
            }
        }
        for r in 0..n {
            let mut s = ZERO;
            for (k, vk) in v.iter().enumerate() {
                s = add(s, mul(h[r * n + col + 1 + k], *vk));
            }
            for (k, vk) in v.iter().enumerate() {
                let idx = r * n + col + 1 + k;
                h[idx] = sub(h[idx], scale(mul(s, conj(*vk)), 2.0));
            }
        }
    }
    let hh = |i: usize, j: usize| h[i * n + j];
    // La Budde: p_k(x) = (x − h_kk) p_{k−1}(x) − Σ_{i=1}^{k−1} h_{k−i,k} (Π β) p_{k−i−1}(x),
    // with β_j = h_{j,j−1}; 1-based k. Polynomials as coefficient vectors.
    let mut polys: Vec<Vec<C>> = vec![vec![ONE]];
    for k in 1..=n {
        let prev = &polys[k - 1];
        let mut p = vec![ZERO; k + 1];
        // x · p_{k−1}
        for (d, c) in prev.iter().enumerate() {
            p[d + 1] = add(p[d + 1], *c);
        }
        // − h_kk · p_{k−1}
        let hkk = hh(k - 1, k - 1);
        for (d, c) in prev.iter().enumerate() {
            p[d] = sub(p[d], mul(hkk, *c));
        }
        let mut prod = ONE;
        for i in 1..k {
            // β_{k−i+1} = h_{(k−i+1),(k−i)} in 1-based terms.
            prod = mul(prod, hh(k - i, k - i - 1));
            let coef = mul(hh(k - i - 1, k - 1), prod);
            for (d, c) in polys[k - i - 1].iter().enumerate() {
                p[d] = sub(p[d], mul(coef, *c));
            }
        }
        polys.push(p);
    }
    // det(xI − M) = Σ_d p[d] x^d, monic; c_j = coefficient of x^{n−j}.
    let p = &polys[n];
    let c = |j: usize| -> C { if j <= n { p[n - j] } else { ZERO } };
    // Newton's identities: p_k = −k c_k − Σ_{i=1}^{k−1} c_i p_{k−i}, c_i = 0 past n.
    let mut traces: Vec<C> = Vec::with_capacity(kmax);
    for k in 1..=kmax {
        let mut s = if k <= n { scale(c(k), -(k as f64)) } else { ZERO };
        for i in 1..k.min(n + 1) {
            s = sub(s, mul(c(i), traces[k - i - 1]));
        }
        traces.push(s);
    }
    traces
}

/// The hafnian of a symmetric `n × n` matrix (`n` even; zero if odd): the sum
/// over perfect matchings of the products of matched entries.
///
/// The `2^{n/2}` subset terms are split into a fixed number of chunks, each
/// summed in order, and the chunk sums are added in order. Threads take whole
/// chunks, so the result is the same bits however many threads run, and the
/// same as a sequential run on a target without threads.
pub fn hafnian(a: &Mat) -> C {
    let n = a.n;
    if n == 0 {
        return ONE;
    }
    if !n.is_multiple_of(2) {
        return ZERO;
    }
    let half = n / 2;
    let total_masks = 1usize << half;
    let chunks = total_masks.min(HAFNIAN_CHUNKS);
    let per = total_masks / chunks;
    let chunk_sum = |c: usize| -> C {
        let mut s = ZERO;
        for mask in c * per..(c + 1) * per {
            s = add(s, hafnian_term(a, half, mask));
        }
        s
    };
    #[cfg(not(target_arch = "wasm32"))]
    let parts: Vec<C> = {
        let threads = std::thread::available_parallelism().map_or(1, |t| t.get()).min(chunks).max(1);
        let mut parts = vec![ZERO; chunks];
        if threads == 1 || half < 12 {
            for (c, p) in parts.iter_mut().enumerate() {
                *p = chunk_sum(c);
            }
        } else {
            std::thread::scope(|sc| {
                let size = chunks.div_ceil(threads);
                for (t, slot) in parts.chunks_mut(size).enumerate() {
                    let chunk_sum = &chunk_sum;
                    sc.spawn(move || {
                        for (k, p) in slot.iter_mut().enumerate() {
                            *p = chunk_sum(t * size + k);
                        }
                    });
                }
            });
        }
        parts
    };
    #[cfg(target_arch = "wasm32")]
    let parts: Vec<C> = (0..chunks).map(chunk_sum).collect();
    parts.into_iter().fold(ZERO, add)
}

/// Subset terms per hafnian are summed in this many fixed chunks.
const HAFNIAN_CHUNKS: usize = 64;

/// One subset's signed term: `(−1)^{n/2−|Z|}` times the `λ^{n/2}` coefficient
/// of `exp(Σ_k tr(Bᵏ) λᵏ / 2k)`, with `B` the rows and columns of `A·X` in the
/// pairs of `Z` (pairs are `(2j, 2j+1)`; `X` swaps each pair).
fn hafnian_term(a: &Mat, half: usize, mask: usize) -> C {
    let pairs: Vec<usize> = (0..half).filter(|j| mask >> j & 1 == 1).collect();
    let m = 2 * pairs.len();
    let mut b = Mat::zeros(m);
    for (r, &pi) in pairs.iter().enumerate() {
        for (cpos, &pj) in pairs.iter().enumerate() {
            for (dr, i) in [2 * pi, 2 * pi + 1].into_iter().enumerate() {
                for (dc, j) in [2 * pj, 2 * pj + 1].into_iter().enumerate() {
                    // (A X)_{ij} = A_{i, j^1}
                    b.a[(2 * r + dr) * m + 2 * cpos + dc] = a.at(i, j ^ 1);
                }
            }
        }
    }
    let traces = power_traces(&b, half);
    let g: Vec<C> = (1..=half).map(|k| scale(traces[k - 1], 1.0 / (2.0 * k as f64))).collect();
    let mut e = vec![ZERO; half + 1];
    e[0] = ONE;
    for k in 1..=half {
        let mut s = ZERO;
        for j in 1..=k {
            s = add(s, scale(mul(g[j - 1], e[k - j]), j as f64));
        }
        e[k] = scale(s, 1.0 / k as f64);
    }
    if (half - pairs.len()).is_multiple_of(2) { e[half] } else { scale(e[half], -1.0) }
}

/// The torontonian of a `2m × 2m` matrix `O` (modes in `(a, a†)` blocks):
/// `Σ_{Z ⊆ [m]} (−1)^{m−|Z|} / √det(I − O_Z)`. Summed in fixed chunks like
/// [`hafnian`], so the bits do not depend on threading.
pub fn torontonian(o: &Mat) -> f64 {
    let m = o.n / 2;
    let total_masks = 1usize << m;
    let chunks = total_masks.min(HAFNIAN_CHUNKS);
    let per = total_masks / chunks;
    let chunk_sum = |c: usize| -> f64 {
        let mut s = 0.0;
        for mask in c * per..(c + 1) * per {
            s += torontonian_term(o, m, mask);
        }
        s
    };
    #[cfg(not(target_arch = "wasm32"))]
    let parts: Vec<f64> = {
        let threads = std::thread::available_parallelism().map_or(1, |t| t.get()).min(chunks).max(1);
        let mut parts = vec![0.0; chunks];
        if threads == 1 || m < 14 {
            for (c, p) in parts.iter_mut().enumerate() {
                *p = chunk_sum(c);
            }
        } else {
            std::thread::scope(|sc| {
                let size = chunks.div_ceil(threads);
                for (t, slot) in parts.chunks_mut(size).enumerate() {
                    let chunk_sum = &chunk_sum;
                    sc.spawn(move || {
                        for (k, p) in slot.iter_mut().enumerate() {
                            *p = chunk_sum(t * size + k);
                        }
                    });
                }
            });
        }
        parts
    };
    #[cfg(target_arch = "wasm32")]
    let parts: Vec<f64> = (0..chunks).map(chunk_sum).collect();
    parts.into_iter().fold(0.0, |a, b| a + b)
}

/// One subset's signed term of the torontonian.
fn torontonian_term(o: &Mat, m: usize, mask: usize) -> f64 {
    let z: Vec<usize> = (0..m).filter(|j| mask >> j & 1 == 1).collect();
    let term = if z.is_empty() {
        1.0
    } else {
        let idx: Vec<usize> = z.iter().copied().chain(z.iter().map(|&j| j + m)).collect();
        let k = idx.len();
        let mut d = o.sub(&idx);
        for i in 0..k {
            for j in 0..k {
                let v = d.a[i * k + j];
                d.a[i * k + j] = if i == j { sub(ONE, v) } else { scale(v, -1.0) };
            }
        }
        1.0 / d.det().re.sqrt()
    };
    if (m - z.len()).is_multiple_of(2) { term } else { -term }
}

/// A Gaussian state of `modes` modes, as its complex covariance matrix in the
/// `(a₁…a_N, a₁†…a_N†)` basis (vacuum = `I/2`), with zero displacement.
#[derive(Clone, Debug)]
pub struct GaussianState {
    pub modes: usize,
    pub cov: Mat,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum GbsError {
    /// Squeezing, interferometer or transmission of the wrong size or range.
    Invalid,
    /// A pattern longer than the state, or a mode index outside it.
    BadPattern,
}

impl GaussianState {
    /// Single-mode squeezed vacua (`r[i]` per mode, 0 = vacuum) through the
    /// interferometer `u` (`N × N` unitary, row-major), then uniform
    /// transmission `eta` on every mode.
    pub fn squeezed(r: &[f64], u: &Mat, eta: f64) -> Result<GaussianState, GbsError> {
        let n = r.len();
        if u.n != n || !(0.0..=1.0).contains(&eta) || r.iter().any(|x| !x.is_finite()) {
            return Err(GbsError::Invalid);
        }
        // Squeezed covariance: per mode [[cosh 2r, −sinh 2r], [−sinh 2r, cosh 2r]] / 2.
        let mut cov = Mat::zeros(2 * n);
        for (i, &ri) in r.iter().enumerate() {
            let e2 = crate::repro::exp(2.0 * ri);
            let (ch, sh) = ((e2 + 1.0 / e2) / 2.0, (e2 - 1.0 / e2) / 2.0);
            cov.set(i, i, C { re: ch / 2.0, im: 0.0 });
            cov.set(n + i, n + i, C { re: ch / 2.0, im: 0.0 });
            cov.set(i, n + i, C { re: -sh / 2.0, im: 0.0 });
            cov.set(n + i, i, C { re: -sh / 2.0, im: 0.0 });
        }
        // Interferometer: S = U ⊕ U*, σ → S σ S†.
        let mut s = Mat::zeros(2 * n);
        for i in 0..n {
            for j in 0..n {
                s.set(i, j, u.at(i, j));
                s.set(n + i, n + j, conj(u.at(i, j)));
            }
        }
        let mut cov = s.mul(&cov).mul(&s.dagger());
        // Loss: σ → η σ + (1 − η) I/2.
        for v in cov.a.iter_mut() {
            *v = scale(*v, eta);
        }
        for i in 0..2 * n {
            cov.a[i * 2 * n + i] = add(cov.a[i * 2 * n + i], C { re: (1.0 - eta) / 2.0, im: 0.0 });
        }
        Ok(GaussianState { modes: n, cov })
    }

    /// The reduced state on `keep` (in that order).
    pub fn reduced(&self, keep: &[usize]) -> GaussianState {
        let n = self.modes;
        let idx: Vec<usize> = keep.iter().copied().chain(keep.iter().map(|&k| k + n)).collect();
        GaussianState { modes: keep.len(), cov: self.cov.sub(&idx) }
    }

    fn q(&self) -> Mat {
        let mut q = self.cov.clone();
        for i in 0..2 * self.modes {
            q.a[i * 2 * self.modes + i] = add(q.a[i * 2 * self.modes + i], C { re: 0.5, im: 0.0 });
        }
        q
    }

    /// Probability that exactly the modes with `pattern[i] = true` click.
    pub fn click_probability(&self, pattern: &[bool]) -> Result<f64, GbsError> {
        if pattern.len() != self.modes {
            return Err(GbsError::BadPattern);
        }
        let q = self.q();
        let qinv = q.inverse().ok_or(GbsError::Invalid)?;
        let n = self.modes;
        let clicked: Vec<usize> = (0..n).filter(|&i| pattern[i]).collect();
        let idx: Vec<usize> = clicked.iter().copied().chain(clicked.iter().map(|&k| k + n)).collect();
        // O = I − Q⁻¹ on the clicked modes.
        let mut o = qinv.sub(&idx);
        let k = idx.len();
        for i in 0..k {
            for j in 0..k {
                let v = o.a[i * k + j];
                o.a[i * k + j] = if i == j { sub(ONE, v) } else { scale(v, -1.0) };
            }
        }
        Ok(torontonian(&o) / q.det().re.sqrt())
    }

    /// Probability of exactly `counts[i]` photons in each mode
    /// (photon-number-resolving detection): `Haf(A_n) / (Π nᵢ! √det Q)`, with
    /// `A = X (I − Q⁻¹)*`.
    pub fn photon_probability(&self, counts: &[usize]) -> Result<f64, GbsError> {
        if counts.len() != self.modes {
            return Err(GbsError::BadPattern);
        }
        let n = self.modes;
        let q = self.q();
        let qinv = q.inverse().ok_or(GbsError::Invalid)?;
        // A = X (I − Q⁻¹)*.
        let mut a = Mat::zeros(2 * n);
        for i in 0..2 * n {
            let xi = if i < n { i + n } else { i - n };
            for j in 0..2 * n {
                let v = conj(qinv.at(xi, j));
                let id = if xi == j { ONE } else { ZERO };
                a.set(i, j, sub(conj(id), v));
            }
        }
        let mut idx = Vec::new();
        for (i, &c) in counts.iter().enumerate() {
            idx.extend(std::iter::repeat_n(i, c));
        }
        let half: Vec<usize> = idx.clone();
        idx.extend(half.iter().map(|&i| i + n));
        let h = hafnian(&a.sub(&idx));
        let mut fact = 1.0;
        for &c in counts {
            for k in 2..=c {
                fact *= k as f64;
            }
        }
        Ok(h.re / (fact * q.det().re.sqrt()))
    }

    /// Exact samples of click patterns, drawn mode by mode: the probability of
    /// each prefix is a torontonian of the reduced state on the modes so far.
    pub fn sample_clicks(&self, shots: usize, seed: u64) -> Vec<Vec<bool>> {
        let mut rng = seed;
        let mut next = || {
            rng = rng.wrapping_add(0x9e37_79b9_7f4a_7c15);
            let mut z = rng;
            z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
            z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
            (((z ^ (z >> 31)) >> 11) as f64 + 0.5) / 9_007_199_254_740_992.0
        };
        let reduced: Vec<GaussianState> = (1..=self.modes).map(|k| self.reduced(&(0..k).collect::<Vec<_>>())).collect();
        let mut out = Vec::with_capacity(shots);
        for _ in 0..shots {
            let mut pattern: Vec<bool> = Vec::with_capacity(self.modes);
            let mut p_prefix = 1.0;
            for st in &reduced {
                pattern.push(true);
                let p_click = st.click_probability(&pattern).unwrap_or(0.0);
                let u = next();
                if u * p_prefix < p_click {
                    p_prefix = p_click;
                } else {
                    *pattern.last_mut().unwrap() = false;
                    p_prefix -= p_click;
                }
            }
            out.push(pattern);
        }
        out
    }

    /// One-mode click probabilities `pᵢ` and two-mode joint click
    /// probabilities `pᵢⱼ` (row-major, `pᵢᵢ = pᵢ`): the low-order statistics
    /// every sampler must reproduce, polynomial in the number of modes.
    pub fn click_marginals(&self) -> (Vec<f64>, Vec<f64>) {
        let n = self.modes;
        let mut p1 = vec![0.0; n];
        for (i, p) in p1.iter_mut().enumerate() {
            *p = self.reduced(&[i]).click_probability(&[true]).unwrap_or(f64::NAN);
        }
        let mut p2 = vec![0.0; n * n];
        for i in 0..n {
            p2[i * n + i] = p1[i];
            for j in i + 1..n {
                let v = self.reduced(&[i, j]).click_probability(&[true, true]).unwrap_or(f64::NAN);
                p2[i * n + j] = v;
                p2[j * n + i] = v;
            }
        }
        (p1, p2)
    }
}

/// How observed click patterns compare with a state's exact low-order
/// statistics.
#[derive(Clone, Debug, PartialEq)]
pub struct MarginalCheck {
    pub shots: usize,
    /// Largest |z| over the one-mode click frequencies.
    pub worst_single: f64,
    /// Largest |z| over the two-mode joint click frequencies.
    pub worst_pair: f64,
    /// Mean clicks per shot, observed and exact.
    pub mean_clicks: (f64, f64),
}

impl GaussianState {
    /// Compare observed click patterns with the exact one- and two-mode click
    /// statistics: a z-score per mode and per pair (binomial standard error),
    /// and the mean number of clicks. A sampler, a device included, that
    /// samples this state reproduces all of them within sampling error.
    pub fn check_samples(&self, samples: &[Vec<bool>]) -> Result<MarginalCheck, GbsError> {
        let n = self.modes;
        if samples.iter().any(|s| s.len() != n) || samples.is_empty() {
            return Err(GbsError::BadPattern);
        }
        let shots = samples.len() as f64;
        let (p1, p2) = self.click_marginals();
        let z = |count: usize, p: f64| -> f64 {
            let f = count as f64 / shots;
            let se = (p * (1.0 - p) / shots).sqrt();
            if se > 0.0 { (f - p) / se } else if f == p { 0.0 } else { f64::INFINITY }
        };
        let mut worst_single: f64 = 0.0;
        for i in 0..n {
            let c = samples.iter().filter(|s| s[i]).count();
            worst_single = worst_single.max(z(c, p1[i]).abs());
        }
        let mut worst_pair: f64 = 0.0;
        for i in 0..n {
            for j in i + 1..n {
                let c = samples.iter().filter(|s| s[i] && s[j]).count();
                worst_pair = worst_pair.max(z(c, p2[i * n + j]).abs());
            }
        }
        let observed = samples.iter().map(|s| s.iter().filter(|&&b| b).count()).sum::<usize>() as f64 / shots;
        Ok(MarginalCheck { shots: samples.len(), worst_single, worst_pair, mean_clicks: (observed, p1.iter().sum()) })
    }
}

/// A Haar-random unitary from a seed (QR of a complex Gaussian matrix, with
/// the phases of R's diagonal fixed), for building test interferometers.
#[allow(clippy::needless_range_loop)]
pub fn random_unitary(n: usize, seed: u64) -> Mat {
    let mut s = seed;
    let mut unit = || {
        s = s.wrapping_add(0x9e37_79b9_7f4a_7c15);
        let mut z = s;
        z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
        z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
        (((z ^ (z >> 31)) >> 11) as f64 + 0.5) / 9_007_199_254_740_992.0
    };
    // Box–Muller with in-crate logarithm and sine/cosine-free polar form.
    let mut gauss = || loop {
        let (x, y) = (2.0 * unit() - 1.0, 2.0 * unit() - 1.0);
        let r2 = x * x + y * y;
        if r2 > 0.0 && r2 < 1.0 {
            let f = (-2.0 * crate::repro::ln(r2) / r2).sqrt();
            return C { re: x * f, im: y * f };
        }
    };
    let mut cols: Vec<Vec<C>> = (0..n).map(|_| (0..n).map(|_| gauss()).collect()).collect();
    for j in 0..n {
        for i in 0..j {
            let mut dot = ZERO;
            for r in 0..n {
                dot = add(dot, mul(conj(cols[i][r]), cols[j][r]));
            }
            for r in 0..n {
                let t = mul(cols[i][r], dot);
                cols[j][r] = sub(cols[j][r], t);
            }
        }
        let norm = cols[j].iter().map(|c| abs2(*c)).sum::<f64>().sqrt();
        for x in cols[j].iter_mut() {
            *x = scale(*x, 1.0 / norm);
        }
    }
    let mut m = Mat::zeros(n);
    for r in 0..n {
        for c in 0..n {
            m.a[r * n + c] = cols[c][r];
        }
    }
    m
}

#[cfg(test)]
mod tests {
    use super::*;

    fn close(a: f64, b: f64, tol: f64) -> bool {
        (a - b).abs() <= tol * (1.0 + b.abs())
    }

    /// Hafnian by exhaustion over perfect matchings.
    fn brute_hafnian(a: &Mat) -> C {
        fn go(a: &Mat, left: &mut Vec<usize>) -> C {
            if left.is_empty() {
                return ONE;
            }
            let i = left.remove(0);
            let mut s = ZERO;
            for k in 0..left.len() {
                let j = left.remove(k);
                s = add(s, mul(a.at(i, j), go(a, left)));
                left.insert(k, j);
            }
            left.insert(0, i);
            s
        }
        go(a, &mut (0..a.n).collect())
    }

    fn random_symmetric(n: usize, seed: u64) -> Mat {
        let u = random_unitary(n, seed);
        let mut m = Mat::zeros(n);
        for i in 0..n {
            for j in i..n {
                let v = add(u.at(i, j), u.at(j, i));
                m.set(i, j, v);
                m.set(j, i, v);
            }
        }
        m
    }

    #[test]
    fn the_hafnian_matches_exhaustion_and_known_values() {
        for n in [2, 4, 6, 8, 10] {
            for seed in 0..5 {
                let a = random_symmetric(n, 100 * n as u64 + seed);
                let (h, b) = (hafnian(&a), brute_hafnian(&a));
                assert!(close(h.re, b.re, 1e-9) && close(h.im, b.im, 1e-9), "n {n} seed {seed}: {h:?} vs {b:?}");
            }
        }
        // haf(J_{2k}) = (2k − 1)!!
        for (n, want) in [(2, 1.0), (4, 3.0), (6, 15.0), (8, 105.0), (10, 945.0), (12, 10395.0)] {
            let j = Mat { n, a: vec![ONE; n * n] };
            assert!(close(hafnian(&j).re, want, 1e-12), "J_{n}");
        }
        assert_eq!(hafnian(&Mat::zeros(3)).re, 0.0);
    }

    #[test]
    fn determinant_and_inverse_agree() {
        let a = random_symmetric(6, 9);
        let inv = a.inverse().unwrap();
        let id = a.mul(&inv);
        for i in 0..6 {
            for j in 0..6 {
                let want = if i == j { 1.0 } else { 0.0 };
                assert!((id.at(i, j).re - want).abs() < 1e-12 && id.at(i, j).im.abs() < 1e-12);
            }
        }
        let (d, di) = (a.det(), inv.det());
        let p = mul(d, di);
        assert!((p.re - 1.0).abs() < 1e-10 && p.im.abs() < 1e-10);
    }

    #[test]
    fn single_mode_squeezed_vacuum_has_the_textbook_statistics() {
        let r = 0.7;
        let st = GaussianState::squeezed(&[r], &Mat::identity(1), 1.0).unwrap();
        let cosh = (crate::repro::exp(r) + crate::repro::exp(-r)) / 2.0;
        let tanh = (crate::repro::exp(2.0 * r) - 1.0) / (crate::repro::exp(2.0 * r) + 1.0);
        assert!(close(st.click_probability(&[true]).unwrap(), 1.0 - 1.0 / cosh, 1e-12));
        assert!(close(st.photon_probability(&[0]).unwrap(), 1.0 / cosh, 1e-12));
        assert!(close(st.photon_probability(&[2]).unwrap(), tanh * tanh / (2.0 * cosh), 1e-12));
        assert!(st.photon_probability(&[1]).unwrap().abs() < 1e-14);
        // P(4) = 4!/(2² 2!)² tanh⁴ / cosh = 3/8 tanh⁴ / cosh.
        assert!(close(st.photon_probability(&[4]).unwrap(), 3.0 / 8.0 * tanh.powi(4) / cosh, 1e-12));
    }

    #[test]
    fn click_probabilities_sum_to_one_with_and_without_loss() {
        for (eta, seed) in [(1.0, 1u64), (0.6, 2), (0.3, 3)] {
            let n = 5;
            let r = [0.9, 0.4, 0.0, 0.7, 0.2];
            let st = GaussianState::squeezed(&r, &random_unitary(n, seed), eta).unwrap();
            let mut total = 0.0;
            for mask in 0..1u32 << n {
                let p: Vec<bool> = (0..n).map(|i| mask >> i & 1 == 1).collect();
                let v = st.click_probability(&p).unwrap();
                assert!(v > -1e-12, "negative probability {v}");
                total += v;
            }
            assert!((total - 1.0).abs() < 1e-10, "eta {eta}: {total}");
        }
    }

    #[test]
    fn photon_probabilities_sum_towards_one() {
        // Three modes, total photons up to a cutoff; the tail is tiny at low squeezing.
        let st = GaussianState::squeezed(&[0.3, 0.2, 0.25], &random_unitary(3, 4), 0.8).unwrap();
        let mut total = 0.0;
        for a in 0..8 {
            for b in 0..8 - a {
                for c in 0..8 - a - b {
                    total += st.photon_probability(&[a, b, c]).unwrap();
                }
            }
        }
        assert!(total > 0.99999 && total < 1.0 + 1e-10, "{total}");
    }

    #[test]
    fn exact_samples_reproduce_the_marginals() {
        let st = GaussianState::squeezed(&[0.8, 0.5, 0.6, 0.3], &random_unitary(4, 7), 0.7).unwrap();
        let shots = 40_000;
        let samples = st.sample_clicks(shots, 11);
        assert_eq!(samples, st.sample_clicks(shots, 11));
        let (p1, p2) = st.click_marginals();
        let n = 4;
        for i in 0..n {
            let f = samples.iter().filter(|s| s[i]).count() as f64 / shots as f64;
            let sigma = (p1[i] * (1.0 - p1[i]) / shots as f64).sqrt();
            assert!((f - p1[i]).abs() < 5.0 * sigma, "mode {i}: {f} vs {}", p1[i]);
            for j in i + 1..n {
                let f = samples.iter().filter(|s| s[i] && s[j]).count() as f64 / shots as f64;
                let q = p2[i * n + j];
                let sigma = (q * (1.0 - q) / shots as f64).sqrt();
                assert!((f - q).abs() < 5.0 * sigma, "modes {i},{j}: {f} vs {q}");
            }
        }
        // Full patterns too: each pattern's frequency against its exact probability.
        for mask in 0..16u32 {
            let p: Vec<bool> = (0..n).map(|i| mask >> i & 1 == 1).collect();
            let want = st.click_probability(&p).unwrap();
            let f = samples.iter().filter(|s| **s == p).count() as f64 / shots as f64;
            let sigma = (want * (1.0 - want) / shots as f64).sqrt().max(1e-4);
            assert!((f - want).abs() < 5.0 * sigma, "pattern {mask:04b}: {f} vs {want}");
        }
    }

    #[test]
    fn exact_samples_pass_the_check_and_a_wrong_state_fails_it() {
        let st = GaussianState::squeezed(&[0.9, 0.7, 0.8, 0.6, 0.5], &random_unitary(5, 21), 0.5).unwrap();
        let samples = st.sample_clicks(20_000, 3);
        let ok = st.check_samples(&samples).unwrap();
        assert!(ok.worst_single < 5.0 && ok.worst_pair < 5.0, "{ok:?}");
        assert!((ok.mean_clicks.0 - ok.mean_clicks.1).abs() < 0.03, "{ok:?}");
        // The same samples against a different interferometer must fail.
        let other = GaussianState::squeezed(&[0.9, 0.7, 0.8, 0.6, 0.5], &random_unitary(5, 22), 0.5).unwrap();
        let bad = other.check_samples(&samples).unwrap();
        assert!(bad.worst_single > 5.0 || bad.worst_pair > 5.0, "{bad:?}");
    }

    #[test]
    fn bad_input_is_refused() {
        assert_eq!(GaussianState::squeezed(&[0.1, 0.2], &Mat::identity(3), 1.0).unwrap_err(), GbsError::Invalid);
        assert_eq!(GaussianState::squeezed(&[0.1], &Mat::identity(1), 1.5).unwrap_err(), GbsError::Invalid);
        let st = GaussianState::squeezed(&[0.1], &Mat::identity(1), 1.0).unwrap();
        assert_eq!(st.click_probability(&[true, false]).unwrap_err(), GbsError::BadPattern);
    }
}