open-standards/wai/wai-quantum/src/quantum_mitigate.rs

//! Deterministic quantum ERROR MITIGATION — `wai.quantum.mitigate`
//! (extensions/quantum-ops § Mitigation).
//!
//! Every quantum error-mitigation technique splits into two halves: **(A)** a
//! hardware data-acquisition front-end you do not own (running a circuit at
//! amplified noise, calibrating a readout matrix, taking randomized snapshots),
//! and **(B)** a deterministic numerical back-end — the extrapolation, the matrix
//! inversion, the shadow estimator — that turns the raw counts into the reported
//! number. Run in float, half (B) is reproducible only on the same build, and
//! nothing records the step between the QPU and the reported number.
//!
//! This module is half (B), done the WAI way: pure fixed-point integer numerics,
//! byte-identical on every machine, each producing a signed, content-addressed,
//! joule-metered [`crate::quantum_ops::MitigationReceipt`] binding
//! `(raw data, noise-model hash, method, seed) → (mitigated estimate ± error)`.
//!
//! Three methods, chosen because they need **no learned noise model** and so are
//! byte-exact reproducible end-to-end today:
//! - [`zne_extrapolate`] — Zero-Noise Extrapolation (linear least-squares and
//!   exact Richardson), the cleanest: pure curve-fitting over measured data.
//! - [`readout_mitigate`] — tensored readout-error inversion (the M3 spirit):
//!   deterministic per-qubit 2×2 assignment-matrix inversion.
//! - [`estimate_pauli`] — classical shadows (random single-qubit Pauli basis):
//!   the inverse measurement channel, median-of-means.
//!
//! Honest boundary, identical to the calibration engines: the **front-end is a
//! deterministic simulator** (we have no QPU), exactly the surface these methods
//! are developed against. What is not approximate is the reproducibility of the
//! post-processing and the receipt over it — the property the float stacks lack.
//!
//! For the heavier estimators (PEC/PEA, tensor-network mitigation) the receipt's
//! `noise_model_hash` is the hook: those estimators consume an externally-learned
//! noise model, and the receipt pins *which* model produced the number without
//! WAI claiming to have learned it.

use crate::quantum_ops::{content_hash, GrantRef, MitigationReceipt};
use ed25519_dalek::SigningKey;

/// Fixed-point fractional bits — the WAI quantum unit, scale `2^20`. Expectation
/// values are signed i64 in `[-MIT_ONE, MIT_ONE]`.
pub const MIT_FRAC: u32 = 20;
/// `1.0` in fixed-point.
pub const MIT_ONE: i64 = 1 << MIT_FRAC;

#[inline]
fn fmul(a: i64, b: i64) -> i64 {
    ((a as i128 * b as i128) >> MIT_FRAC) as i64
}

#[inline]
fn fdiv(p: i64, q: i64) -> i64 {
    if q == 0 {
        return 0;
    }
    (((p as i128) << MIT_FRAC) / q as i128) as i64
}

/// Integer square root of a non-negative i128 (floor). Deterministic.
fn isqrt_i128(v: i128) -> i128 {
    if v <= 0 {
        return 0;
    }
    let mut x = v;
    let mut y = (x + 1) / 2;
    while y < x {
        x = y;
        y = (x + v / x) / 2;
    }
    x
}

/// `sqrt` of a fixed-point *square*: given `s = real² · 2^(2·MIT_FRAC)`, returns
/// `real · 2^MIT_FRAC`. (isqrt of a 2·FRAC-scaled square yields a FRAC-scaled
/// value directly.)
fn sqrt_fx_of_square(sq: i128) -> i64 {
    isqrt_i128(sq) as i64
}

fn splitmix64(state: &mut u64) -> u64 {
    *state = state.wrapping_add(0x9E37_79B9_7F4A_7C15);
    let mut z = *state;
    z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
    z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
    z ^ (z >> 31)
}

/// Draw one ±1 outcome with `P(+1) = (1 + e_fx)/2` for a signed expectation
/// `e_fx ∈ [-MIT_ONE, MIT_ONE]`. Deterministic given the stream.
fn draw_pm1(st: &mut u64, e_fx: i64) -> i64 {
    let p_plus = (MIT_ONE + e_fx) / 2; // [0, MIT_ONE]
    let r = (splitmix64(st) >> (64 - MIT_FRAC)) as i64; // [0, MIT_ONE)
    if r < p_plus {
        1
    } else {
        -1
    }
}

// ===========================================================================
// 1. Zero-Noise Extrapolation (wai.quantum.mitigate / zne.*)
// ===========================================================================

/// Which extrapolator to fit to the noise-scaled expectation table.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ZneMethod {
    /// Least-squares straight line `E(λ)=a+bλ`, extrapolated to `λ=0`. Robust to
    /// shot noise, biased when the true decay is curved.
    Linear,
    /// Exact polynomial through all points (Lagrange at `x=0`) — classic
    /// Richardson extrapolation. Unbiased for a polynomial decay, noisier.
    Richardson,
}

impl ZneMethod {
    fn tag(self) -> &'static str {
        match self {
            ZneMethod::Linear => "zne.linear",
            ZneMethod::Richardson => "zne.richardson",
        }
    }
}

/// Output of a ZNE post-processing step. Deterministic given the input table.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ZneResult {
    pub method: &'static str,
    /// The measured `(noise_scale_fx, expectation_fx)` table that went in.
    pub points: Vec<(i64, i64)>,
    /// Expectation at the smallest scale (base hardware noise).
    pub raw_fx: i64,
    /// Extrapolation to zero noise — the mitigated estimate.
    pub mitigated_fx: i64,
    /// Uncertainty: fit RMS residual (Linear) or |Richardson − linear| model
    /// disagreement (Richardson).
    pub error_fx: i64,
    pub shots: u64,
}

/// Least-squares line through `points`, extrapolated to `λ=0`. Returns
/// `(intercept_fx, slope_fx)`.
fn fit_linear(points: &[(i64, i64)]) -> (i64, i64) {
    let n = points.len() as i128;
    if n == 0 {
        return (0, 0);
    }
    let (mut sx, mut sy, mut sxy, mut sxx) = (0i128, 0i128, 0i128, 0i128);
    for &(x, y) in points {
        sx += x as i128;
        sy += y as i128;
        sxy += x as i128 * y as i128; // scale 2^40
        sxx += x as i128 * x as i128; // scale 2^40
    }
    let denom = n * sxx - sx * sx; // scale 2^40
    if denom == 0 {
        return ((sy / n) as i64, 0); // degenerate: flat mean
    }
    let numer = n * sxy - sx * sy; // scale 2^40
    // slope_fx = (numer/denom) · 2^20
    let slope = ((numer << MIT_FRAC) / denom) as i64;
    let mean_x = (sx / n) as i64;
    let mean_y = (sy / n) as i64;
    let intercept = mean_y - fmul(slope, mean_x);
    (intercept, slope)
}

/// Exact Richardson extrapolation (Lagrange interpolation evaluated at `x=0`).
/// Requires distinct scales; capped at 8 points to keep the products bounded.
fn richardson_zero(points: &[(i64, i64)]) -> i64 {
    let m = points.len();
    let mut acc = 0i64;
    for i in 0..m {
        let xi = points[i].0;
        // weight_i = Π_{j≠i} x_j / (x_j − x_i)
        let mut w = MIT_ONE;
        for (j, &(xj, _)) in points.iter().enumerate() {
            if j == i {
                continue;
            }
            let d = xj - xi;
            if d == 0 {
                continue; // duplicate scale: skip (shouldn't happen)
            }
            w = fdiv(fmul(w, xj), d);
        }
        acc += fmul(w, points[i].1);
    }
    acc
}

/// Extrapolate a noise-scaled expectation table to the zero-noise limit. `points`
/// is `(noise_scale_fx, measured_expectation_fx)`, scale `≥ MIT_ONE` (1.0 = base
/// noise). Pure post-processing — no noise model required.
pub fn zne_extrapolate(points: &[(i64, i64)], method: ZneMethod, shots: u64) -> ZneResult {
    let mut pts = points.to_vec();
    pts.sort_by_key(|p| p.0);
    let raw_fx = pts.first().map(|p| p.1).unwrap_or(0);

    let (intercept, slope) = fit_linear(&pts);
    let lin0 = intercept;

    let (mitigated_fx, error_fx) = match method {
        ZneMethod::Linear => {
            // RMS residual against the fitted line.
            let n = pts.len().max(1) as i128;
            let mut ss = 0i128;
            for &(x, y) in &pts {
                let pred = intercept + fmul(slope, x);
                let r = (y - pred) as i128;
                ss += r * r; // scale 2^40
            }
            let rms = sqrt_fx_of_square(ss / n);
            (lin0, rms)
        }
        ZneMethod::Richardson => {
            let capped: Vec<(i64, i64)> = pts.iter().take(8).copied().collect();
            let rich = richardson_zero(&capped);
            (rich, (rich - lin0).abs())
        }
    };

    ZneResult {
        method: method.tag(),
        points: pts,
        raw_fx,
        mitigated_fx,
        error_fx,
        shots,
    }
}

impl ZneResult {
    /// Canonical bytes of the raw input table (what `input_hash` binds).
    pub fn input_bytes(&self) -> Vec<u8> {
        let mut b = Vec::with_capacity(self.points.len() * 16 + 8);
        b.extend_from_slice(b"wai:zne-input\x01");
        for &(x, y) in &self.points {
            b.extend_from_slice(&x.to_le_bytes());
            b.extend_from_slice(&y.to_le_bytes());
        }
        b
    }

    /// Seal this ZNE step into a signed, joule-metered mitigation receipt. No
    /// noise model — the receipt records that none was needed.
    #[allow(clippy::too_many_arguments)]
    pub fn seal(
        &self,
        signer: &SigningKey,
        signer_id: impl Into<String>,
        backend_id: impl Into<String>,
        observable: impl Into<String>,
        joules_micro: u64,
        grant: GrantRef,
    ) -> MitigationReceipt {
        MitigationReceipt::seal(
            signer,
            signer_id,
            backend_id,
            self.method,
            observable,
            content_hash(&self.input_bytes()),
            None,
            self.raw_fx,
            self.mitigated_fx,
            self.error_fx,
            self.shots,
            joules_micro,
            grant,
            None,
        )
    }
}

/// Deterministic front-end for ZNE: a noisy observable whose expectation decays
/// as `E(λ) = ideal · exp(-γλ)` (a depolarizing-like model), sampled with shot
/// noise at each scale. Returns the `(scale_fx, measured_fx)` table a real ZNE
/// run would hand to [`zne_extrapolate`].
pub fn simulate_zne(
    ideal_fx: i64,
    gamma_fx: i64,
    scales_fx: &[i64],
    shots: u32,
    seed: u64,
) -> Vec<(i64, i64)> {
    scales_fx
        .iter()
        .map(|&lam| {
            let e_true = fmul(ideal_fx, exp_neg_fx(fmul(gamma_fx, lam)));
            let mut st = seed
                .wrapping_mul(0x1000_0001)
                .wrapping_add(lam as u64)
                .wrapping_add(0xA5A5);
            let mut sum = 0i64;
            for _ in 0..shots {
                sum += draw_pm1(&mut st, e_true);
            }
            let meas = ((sum as i128 * MIT_ONE as i128) / shots.max(1) as i128) as i64;
            (lam, meas)
        })
        .collect()
}

/// Fixed-point `exp(-x)` for `x ≥ 0` via a 6-term Taylor series with range
/// reduction by halving (`exp(-x) = (exp(-x/2))²`). Deterministic, integer-only.
fn exp_neg_fx(x_fx: i64) -> i64 {
    if x_fx <= 0 {
        return MIT_ONE;
    }
    // Range-reduce so the argument is small, then square back.
    let mut halvings = 0;
    let mut x = x_fx;
    while x > (MIT_ONE >> 2) {
        x >>= 1;
        halvings += 1;
    }
    // Taylor: 1 - x + x²/2 - x³/6 + x⁴/24 - x⁵/120
    let mut term = MIT_ONE;
    let mut acc = 0i64;
    let mut sign = 1i64;
    for k in 0..6 {
        acc += sign * term;
        term = fmul(term, x) / (k + 1) as i64;
        sign = -sign;
    }
    let mut r = acc;
    for _ in 0..halvings {
        r = fmul(r, r);
    }
    r.clamp(0, MIT_ONE)
}

// ===========================================================================
// 2. Readout-error mitigation — tensored inversion (wai.quantum.mitigate / readout.tensored)
// ===========================================================================

/// Per-qubit readout assignment matrices. Each is `[p00, p01, p10, p11]`
/// (row = measured, col = true): `p00 = P(measure 0 | true 0)`,
/// `p01 = P(measure 0 | true 1)`, etc. Fixed-point.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ReadoutModel {
    pub qubits: Vec<[i64; 4]>,
}

impl ReadoutModel {
    /// A symmetric model: each qubit flips `0→1` with `p01_hi` and `1→0` with
    /// `p10_lo` (as fixed-point rates); the rest is correct assignment.
    pub fn symmetric(n: usize, p_meas0_given1: i64, p_meas1_given0: i64) -> ReadoutModel {
        let mat = [
            MIT_ONE - p_meas1_given0, // p00
            p_meas0_given1,           // p01
            p_meas1_given0,           // p10
            MIT_ONE - p_meas0_given1, // p11
        ];
        ReadoutModel {
            qubits: vec![mat; n],
        }
    }

    /// Canonical bytes for content-hashing (the pinned noise model).
    pub fn bytes(&self) -> Vec<u8> {
        let mut b = Vec::with_capacity(self.qubits.len() * 32 + 8);
        b.extend_from_slice(b"wai:readout-model\x01");
        for q in &self.qubits {
            for v in q {
                b.extend_from_slice(&v.to_le_bytes());
            }
        }
        b
    }

    fn inv(&self, q: usize) -> [i64; 4] {
        let [a, b, c, d] = self.qubits[q];
        let det = fmul(a, d) - fmul(b, c);
        let id = fdiv(MIT_ONE, det);
        [fmul(id, d), -fmul(id, b), -fmul(id, c), fmul(id, a)]
    }
}

/// Apply a per-qubit 2×2 operator across bit `q` of a dense length-`2^n` vector
/// (the butterfly). `op = [m00, m01, m10, m11]`.
fn apply_qubit(v: &mut [i64], n: usize, q: usize, op: [i64; 4]) {
    let bit = 1usize << q;
    for x in 0..(1usize << n) {
        if x & bit == 0 {
            let v0 = v[x];
            let v1 = v[x | bit];
            v[x] = fmul(op[0], v0) + fmul(op[1], v1);
            v[x | bit] = fmul(op[2], v0) + fmul(op[3], v1);
        }
    }
}

/// Output of a readout-mitigation step.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ReadoutResult {
    pub n: usize,
    /// Raw expectation of the observable over the measured distribution.
    pub raw_fx: i64,
    /// Mitigated expectation after applying `A^{-1}`.
    pub mitigated_fx: i64,
    /// Negativity of the corrected quasi-distribution (∑ of negative mass) — a
    /// standard diagnostic; 0 means the correction stayed physical.
    pub negativity_fx: i64,
    pub shots: u64,
    /// Canonical measured-counts bytes (what `input_hash` binds).
    pub counts_bytes: Vec<u8>,
    /// Content hash of the readout model (the pinned noise model).
    pub model_hash: [u8; 32],
}

/// Invert tensored readout error. `counts` is measured shots over the `2^n`
/// bitstrings (little-endian: bit `q` is qubit `q`). `obs` is the diagonal
/// observable value per bitstring (fixed-point, e.g. `±MIT_ONE` for a parity).
pub fn readout_mitigate(
    counts: &[u64],
    obs: &[i64],
    model: &ReadoutModel,
) -> ReadoutResult {
    let n = model.qubits.len();
    let dim = 1usize << n;
    assert_eq!(counts.len(), dim, "counts must have 2^n entries");
    assert_eq!(obs.len(), dim, "obs must have 2^n entries");

    let total: u64 = counts.iter().sum::<u64>().max(1);
    let mut p: Vec<i64> = counts
        .iter()
        .map(|&c| ((c as i128 * MIT_ONE as i128) / total as i128) as i64)
        .collect();

    let raw_fx = p
        .iter()
        .zip(obs)
        .map(|(&pi, &oi)| fmul(pi, oi))
        .sum::<i64>();

    for q in 0..n {
        apply_qubit(&mut p, n, q, model.inv(q));
    }

    let mitigated_fx = p.iter().zip(obs).map(|(&pi, &oi)| fmul(pi, oi)).sum::<i64>();
    let negativity_fx = p.iter().filter(|&&x| x < 0).map(|&x| -x).sum::<i64>();

    let mut cb = Vec::with_capacity(dim * 8 + 8);
    cb.extend_from_slice(b"wai:readout-counts\x01");
    for &c in counts {
        cb.extend_from_slice(&c.to_le_bytes());
    }

    ReadoutResult {
        n,
        raw_fx,
        mitigated_fx,
        negativity_fx,
        shots: total,
        counts_bytes: cb,
        model_hash: content_hash(&model.bytes()),
    }
}

impl ReadoutResult {
    /// Seal this readout-mitigation step, binding the pinned readout model.
    pub fn seal(
        &self,
        signer: &SigningKey,
        signer_id: impl Into<String>,
        backend_id: impl Into<String>,
        observable: impl Into<String>,
        joules_micro: u64,
        grant: GrantRef,
    ) -> MitigationReceipt {
        MitigationReceipt::seal(
            signer,
            signer_id,
            backend_id,
            "readout.tensored",
            observable,
            content_hash(&self.counts_bytes),
            Some(self.model_hash),
            self.raw_fx,
            self.mitigated_fx,
            self.negativity_fx,
            self.shots,
            joules_micro,
            grant,
            None,
        )
    }
}

/// Deterministic front-end for readout mitigation: sample measured counts from
/// `A · p_true` where `p_true` is an ideal distribution over `2^n` bitstrings
/// (fixed-point, summing to `MIT_ONE`). Returns the measured counts a real run
/// would produce.
pub fn simulate_readout(
    p_true_fx: &[i64],
    model: &ReadoutModel,
    shots: u32,
    seed: u64,
) -> Vec<u64> {
    let n = model.qubits.len();
    let dim = 1usize << n;
    // Forward-apply A to get the measured distribution.
    let mut p = p_true_fx.to_vec();
    for q in 0..n {
        apply_qubit(&mut p, n, q, model.qubits[q]);
    }
    // Build a CDF (clamp tiny negatives from rounding) and sample.
    let mut cdf = vec![0i64; dim];
    let mut run = 0i64;
    for i in 0..dim {
        run += p[i].max(0);
        cdf[i] = run;
    }
    let norm = run.max(1);
    let mut counts = vec![0u64; dim];
    let mut st = seed.wrapping_mul(0x2545_F491).wrapping_add(1);
    for _ in 0..shots {
        let r = (((splitmix64(&mut st) >> (64 - MIT_FRAC)) as i128 * norm as i128) >> MIT_FRAC) as i64;
        let mut idx = dim - 1;
        for (i, &c) in cdf.iter().enumerate() {
            if r < c {
                idx = i;
                break;
            }
        }
        counts[idx] += 1;
    }
    counts
}

// ===========================================================================
// 3. Classical shadows — random single-qubit Pauli (wai.quantum.mitigate / shadow.pauli)
// ===========================================================================

/// Pauli basis / operator labels used by the shadow estimator.
pub const PX: u8 = 0;
pub const PY: u8 = 1;
pub const PZ: u8 = 2;
/// Identity — a qubit not in the observable's support.
pub const PI: u8 = 3;

/// One classical-shadow snapshot: the randomly chosen measurement basis per
/// qubit (`PX`/`PY`/`PZ`) and the ±1 outcome observed.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Snapshot {
    pub bases: Vec<u8>,
    pub outcomes: Vec<i8>,
}

/// Output of a classical-shadows Pauli-observable estimate.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ShadowResult {
    pub estimate_fx: i64,
    /// Median-of-means spread across batches (half the batch min–max range).
    pub error_fx: i64,
    pub n_snapshots: usize,
    /// Canonical snapshot bytes (what `input_hash` binds).
    pub input: Vec<u8>,
}

fn canonical_snapshots(snaps: &[Snapshot]) -> Vec<u8> {
    let mut b = Vec::with_capacity(snaps.len() * 8 + 8);
    b.extend_from_slice(b"wai:shadow-snapshots\x01");
    for s in snaps {
        for &x in &s.bases {
            b.push(x);
        }
        for &o in &s.outcomes {
            b.push(o as u8);
        }
    }
    b
}

/// Single-snapshot contribution to `⟨P⟩`: for each qubit in `pauli`'s support the
/// measured basis must match `pauli`'s basis (factor `3·outcome`), else the whole
/// snapshot contributes 0. Identity qubits contribute factor 1. Returns the
/// contribution as an exact integer (product of `±3`s).
fn snapshot_contribution(snap: &Snapshot, pauli: &[u8]) -> Option<i64> {
    let mut prod = 1i64;
    for (q, &pq) in pauli.iter().enumerate() {
        if pq == PI {
            continue;
        }
        if snap.bases[q] != pq {
            return None; // wrong basis measured on a support qubit → 0
        }
        prod *= 3 * snap.outcomes[q] as i64;
    }
    Some(prod)
}

/// Estimate `⟨P⟩` for a Pauli string `pauli` (`PX`/`PY`/`PZ`/`PI` per qubit) from
/// classical-shadow snapshots, using a median-of-means over `n_batches` groups
/// for robustness. Pure post-processing; no noise model.
pub fn estimate_pauli(snaps: &[Snapshot], pauli: &[u8], n_batches: usize) -> ShadowResult {
    let m = snaps.len();
    let batches = n_batches.clamp(1, m.max(1));
    let mut means: Vec<i64> = Vec::with_capacity(batches);
    for b in 0..batches {
        let lo = b * m / batches;
        let hi = (b + 1) * m / batches;
        if hi <= lo {
            continue;
        }
        let mut acc = 0i128;
        for snap in &snaps[lo..hi] {
            // contribution·MIT_ONE, 0 when the snapshot misses the basis.
            let c = snapshot_contribution(snap, pauli).unwrap_or(0);
            acc += c as i128 * MIT_ONE as i128;
        }
        means.push((acc / (hi - lo) as i128) as i64);
    }
    means.sort_unstable();
    let estimate_fx = if means.is_empty() {
        0
    } else {
        means[means.len() / 2]
    };
    let error_fx = if means.len() < 2 {
        0
    } else {
        (means[means.len() - 1] - means[0]) / 2
    };
    ShadowResult {
        estimate_fx,
        error_fx,
        n_snapshots: m,
        input: canonical_snapshots(snaps),
    }
}

impl ShadowResult {
    /// Seal this classical-shadows estimate. No noise model — the snapshots and
    /// the estimator are the whole story.
    #[allow(clippy::too_many_arguments)]
    pub fn seal(
        &self,
        signer: &SigningKey,
        signer_id: impl Into<String>,
        backend_id: impl Into<String>,
        observable: impl Into<String>,
        raw_fx: i64,
        joules_micro: u64,
        grant: GrantRef,
    ) -> MitigationReceipt {
        MitigationReceipt::seal(
            signer,
            signer_id,
            backend_id,
            "shadow.pauli",
            observable,
            content_hash(&self.input),
            None,
            raw_fx,
            self.estimate_fx,
            self.error_fx,
            self.n_snapshots as u64,
            joules_micro,
            grant,
            None,
        )
    }
}

/// Deterministic front-end for classical shadows: a product state with per-qubit
/// Bloch vector `[rx, ry, rz]` (fixed-point, `|r| ≤ 1`). For each snapshot each
/// qubit is measured in a uniformly-random Pauli basis, with the correct ±1
/// outcome probabilities. The true `⟨Pauli string⟩` is the product of the Bloch
/// components over the support (see [`bloch_expectation`]) — a known ground truth.
pub fn simulate_shadows(bloch: &[[i64; 3]], n_snapshots: usize, seed: u64) -> Vec<Snapshot> {
    let n = bloch.len();
    let mut out = Vec::with_capacity(n_snapshots);
    let mut st = seed.wrapping_mul(0x9E37_79B9).wrapping_add(7);
    for _ in 0..n_snapshots {
        let mut bases = Vec::with_capacity(n);
        let mut outcomes = Vec::with_capacity(n);
        for qb in bloch.iter() {
            let basis = (splitmix64(&mut st) % 3) as u8; // 0=X,1=Y,2=Z uniform
            let r = qb[basis as usize];
            let o = draw_pm1(&mut st, r) as i8;
            bases.push(basis);
            outcomes.push(o);
        }
        out.push(Snapshot { bases, outcomes });
    }
    out
}

/// The analytic `⟨Pauli string⟩` for a Bloch product state — the ground truth the
/// shadow estimate should recover. Product of the Bloch component along each
/// support qubit's Pauli axis.
pub fn bloch_expectation(bloch: &[[i64; 3]], pauli: &[u8]) -> i64 {
    let mut prod = MIT_ONE;
    for (q, &pq) in pauli.iter().enumerate() {
        if pq == PI {
            continue;
        }
        prod = fmul(prod, bloch[q][pq as usize]);
    }
    prod
}

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

    fn key(s: u8) -> SigningKey {
        SigningKey::from_bytes(&[s; 32])
    }

    // --- ZNE ----------------------------------------------------------------

    #[test]
    fn zne_recovers_ideal_from_noisy_decay() {
        // ideal ⟨O⟩ = 0.8, exponential decay in the noise scale.
        let ideal = (0.8 * MIT_ONE as f64) as i64;
        let gamma = (0.25 * MIT_ONE as f64) as i64;
        // Canonical, well-conditioned ZNE: 3 scales → quadratic Richardson. (More
        // points raise the polynomial degree and reintroduce Runge instability
        // when extrapolating back to λ=0 — the classic ZNE bias/variance trade.)
        let scales: Vec<i64> = [1.0, 2.0, 3.0]
            .iter()
            .map(|s| (s * MIT_ONE as f64) as i64)
            .collect();
        let table = simulate_zne(ideal, gamma, &scales, 20_000, 0xBEEF);
        let raw = table[0].1;

        let rich = zne_extrapolate(&table, ZneMethod::Richardson, 100_000);
        // Richardson should pull the estimate much closer to the ideal than the
        // raw base-noise value.
        let err_raw = (raw - ideal).abs();
        let err_mit = (rich.mitigated_fx - ideal).abs();
        assert!(
            err_mit < err_raw,
            "mitigated must beat raw: raw_err={err_raw}, mit_err={err_mit}"
        );
        // within ~8% of ideal
        assert!(err_mit < (0.08 * MIT_ONE as f64) as i64, "mit within 8%: {err_mit}");
    }

    #[test]
    fn zne_is_deterministic() {
        let ideal = MIT_ONE / 2;
        let g = MIT_ONE / 5;
        let s: Vec<i64> = [1.0, 2.0, 3.0].iter().map(|x| (x * MIT_ONE as f64) as i64).collect();
        let a = simulate_zne(ideal, g, &s, 5000, 42);
        let b = simulate_zne(ideal, g, &s, 5000, 42);
        assert_eq!(a, b, "same seed → byte-identical table");
        assert_eq!(
            zne_extrapolate(&a, ZneMethod::Linear, 15_000),
            zne_extrapolate(&b, ZneMethod::Linear, 15_000)
        );
    }

    #[test]
    fn zne_linear_exact_on_a_line() {
        // On a perfectly linear table the linear fit is exact and residual 0.
        let pts = vec![
            (MIT_ONE, MIT_ONE / 2),         // (1.0, 0.5)
            (2 * MIT_ONE, MIT_ONE / 4),     // (2.0, 0.25) slope -0.25
            (3 * MIT_ONE, 0),               // (3.0, 0.0)
        ];
        let r = zne_extrapolate(&pts, ZneMethod::Linear, 0);
        // extrapolate to 0 → 0.75
        assert!((r.mitigated_fx - (3 * MIT_ONE / 4)).abs() < 16, "intercept ~0.75");
        assert!(r.error_fx < 16, "residual ~0 on an exact line");
    }

    #[test]
    fn zne_seals_and_verifies_no_model() {
        let pts = vec![(MIT_ONE, MIT_ONE / 2), (2 * MIT_ONE, MIT_ONE / 4)];
        let r = zne_extrapolate(&pts, ZneMethod::Richardson, 10_000);
        let rec = r.seal(&key(1), "did:key:lab", "wai.quantum.circuit", "ZZ", 500_000, GrantRef::unbounded("quantum.mitigate"));
        assert!(rec.verify());
        assert!(rec.noise_model_hash.is_none());
        assert!(rec.input_matches(&r.input_bytes()));
    }

    // --- Readout ------------------------------------------------------------

    #[test]
    fn readout_inversion_recovers_parity() {
        let n = 2;
        // GHZ-like ideal: |00> and |11> each 0.5. Parity ZZ = +1 on both → ⟨ZZ⟩=1.
        let dim = 1usize << n;
        let mut p_true = vec![0i64; dim];
        p_true[0b00] = MIT_ONE / 2;
        p_true[0b11] = MIT_ONE / 2;
        let obs: Vec<i64> = (0..dim)
            .map(|x| if (x.count_ones() & 1) == 0 { MIT_ONE } else { -MIT_ONE })
            .collect();

        let model = ReadoutModel::symmetric(n, MIT_ONE / 12, MIT_ONE / 12); // ~8% each way
        let counts = simulate_readout(&p_true, &model, 200_000, 0xC0FFEE);
        let res = readout_mitigate(&counts, &obs, &model);

        // Raw parity is pulled below 1 by readout error; mitigation restores it.
        assert!(res.raw_fx < (0.9 * MIT_ONE as f64) as i64, "raw pulled down: {}", res.raw_fx);
        let err = (res.mitigated_fx - MIT_ONE).abs();
        assert!(err < (0.03 * MIT_ONE as f64) as i64, "mitigated ~1.0: {}", res.mitigated_fx);
    }

    #[test]
    fn readout_seals_with_pinned_model() {
        let n = 2;
        let dim = 1usize << n;
        let mut p_true = vec![0i64; dim];
        p_true[0] = MIT_ONE;
        let obs = vec![MIT_ONE, -MIT_ONE, -MIT_ONE, MIT_ONE];
        let model = ReadoutModel::symmetric(n, MIT_ONE / 20, MIT_ONE / 20);
        let counts = simulate_readout(&p_true, &model, 50_000, 1);
        let res = readout_mitigate(&counts, &obs, &model);
        let rec = res.seal(&key(2), "lab", "sim", "ZZ", 400_000, GrantRef::unbounded("quantum.mitigate"));
        assert!(rec.verify());
        assert!(rec.noise_model_matches(&model.bytes()));
        assert!(rec.input_matches(&res.counts_bytes));
    }

    // --- Classical shadows --------------------------------------------------

    #[test]
    fn shadows_recover_pauli_expectation() {
        // 3-qubit product state, known Bloch vectors.
        let bloch = [
            [(0.6 * MIT_ONE as f64) as i64, 0, (0.5 * MIT_ONE as f64) as i64],
            [0, (0.7 * MIT_ONE as f64) as i64, 0],
            [0, 0, (0.9 * MIT_ONE as f64) as i64],
        ];
        // Observable X0 Y1 Z2 → truth = 0.6 · 0.7 · 0.9.
        let pauli = [PX, PY, PZ];
        let truth = bloch_expectation(&bloch, &pauli);
        let snaps = simulate_shadows(&bloch, 60_000, 0x5EED);
        let est = estimate_pauli(&snaps, &pauli, 12);
        let err = (est.estimate_fx - truth).abs();
        assert!(
            err < (0.05 * MIT_ONE as f64) as i64,
            "shadow estimate {} vs truth {} (err {})",
            est.estimate_fx, truth, err
        );
    }

    #[test]
    fn shadows_deterministic_and_sealed() {
        let bloch = [[0, 0, MIT_ONE], [0, 0, MIT_ONE / 2]];
        let a = simulate_shadows(&bloch, 4000, 9);
        let b = simulate_shadows(&bloch, 4000, 9);
        assert_eq!(a, b);
        let pauli = [PZ, PI];
        let ea = estimate_pauli(&a, &pauli, 8);
        assert_eq!(ea, estimate_pauli(&b, &pauli, 8));
        let rec = ea.seal(&key(3), "lab", "sim", "Z0", 0, 300_000, GrantRef::unbounded("quantum.mitigate"));
        assert!(rec.verify());
        assert!(rec.input_matches(&ea.input));
    }

    #[test]
    fn shadow_wrong_basis_contributes_zero() {
        // A snapshot that measured Z on qubit 0 cannot inform ⟨X0⟩.
        let snap = Snapshot { bases: vec![PZ], outcomes: vec![1] };
        assert_eq!(snapshot_contribution(&snap, &[PX]), None);
        let snap2 = Snapshot { bases: vec![PX], outcomes: vec![-1] };
        assert_eq!(snapshot_contribution(&snap2, &[PX]), Some(-3));
    }
}