wai-quantum 0.3.19

//! Deterministic FILTER-FUNCTION control + NOISE SPECTROSCOPY —
//! `wai.quantum.cal.spectroscopy` + `wai.quantum.cal.robust_ff`
//! (extensions/quantum-ops § Calibration, frequency domain).
//!
//! The composite-pulse robust control already in [`crate::quantum_cal`] hardens a
//! gate against a *static* amplitude error. Real devices are limited by
//! *correlated, low-frequency* noise — 1/f dephasing above all — and the SOTA
//! tools address it in the frequency domain via the **filter function**: a
//! control's first-order sensitivity `F(ω)` to dephasing noise, so the coherence
//! decay is the overlap `χ = (1/π) ∫ S(ω) F(ω)/ω² dω` of the noise PSD `S(ω)`
//! with the filter. Two matched capabilities fall straight out of that one
//! integral, and both run **entirely offline against a model** — no hardware
//! front-end at all, so they are 100% deterministic-Rust and 100% receiptable:
//!
//! - **Noise spectroscopy** ([`reconstruct_psd`]) — a bank of dynamical-decoupling
//!   sequences are narrow-band filters at different centre frequencies; the
//!   coherence they each retain measures `S(ω)` there. A Richardson–Lucy
//!   deconvolution over the filter bank reconstructs the hidden PSD.
//! - **Filter-function robust control** ([`optimize_dd`]) — given a PSD, reshape a
//!   DD sequence's pulse timing to push its filter's stopband over where the noise
//!   lives, minimising `χ` (maximising coherence). The optimiser reshapes CPMG
//!   toward the noise-avoiding sequence for that spectrum.
//!
//! The matched pair: spectroscopy signs a PSD; robust control signs a pulse whose
//! filter provably avoids that PSD, binding the PSD's hash. Everything is pure
//! fixed-point integer arithmetic (reusing the calibration floor's `sin`/`cos`/
//! `exp`), byte-identical on every machine, and seals a `wai.quantum.calibration`
//! receipt via [`crate::quantum_cal::seal_artifacts`].
//!
//! Honest boundary: the qubit response is a clean first-order dephasing model, and
//! the "measurement" is a deterministic simulator (no hardware). What is not
//! approximate is the reproducibility of the DSP and the receipt over it.

use crate::quantum_cal::{cos_fx, exp_neg_fx, seal_artifacts, sin_fx, CalArtifacts, CAL_FRAC, CAL_ONE};
use crate::quantum_ops::{content_hash, CalibrationReceipt, GrantRef};
use ed25519_dalek::SigningKey;

const ONE: i64 = CAL_ONE;
const FRAC: u32 = CAL_FRAC;
const PI_FX: i64 = 3_294_199; // round(π · 2^20)
const INV_PI_FX: i64 = 333_772; // round(2^20 / π)

#[inline]
fn fmul(a: i64, b: i64) -> i64 {
    ((a as i128 * b as i128) >> FRAC) as i64
}
#[inline]
fn fdiv(a: i64, b: i64) -> i64 {
    if b == 0 {
        return 0;
    }
    (((a as i128) << FRAC) / b as i128) as i64
}

/// Fixed-point natural log for `x ∈ (0, ONE]`, via `ln x = 2·artanh((x−1)/(x+1))`.
/// Returns a signed value (≤ 0 for `x ≤ 1`). Deterministic.
fn ln_fx(x: i64) -> i64 {
    if x <= 0 {
        return -20 * ONE; // guard: treat as a large negative log
    }
    let z = fdiv(x - ONE, x + ONE); // (x-1)/(x+1) ∈ (-1, 0]
    let z2 = fmul(z, z);
    let mut term = z;
    let mut acc = 0i64;
    let mut k = 1i64;
    for _ in 0..12 {
        acc += term / k;
        term = fmul(term, z2);
        k += 2;
    }
    2 * acc
}

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)
}

// ===========================================================================
// Frequency grid + noise PSD model
// ===========================================================================

/// A discrete angular-frequency grid (fixed-point ω values + spacing). The pulse
/// interval is normalised to `T = 1` (`ONE`), so ω is in units of `1/T`.
#[derive(Clone, Debug)]
pub struct NoiseGrid {
    pub omega: Vec<i64>,
    pub dw: i64,
}

impl NoiseGrid {
    /// Linear grid of `n` points over `[w_min, w_max]` (given as f64 for setup;
    /// the grid values themselves are fixed-point and drive all later arithmetic).
    pub fn linear(w_min: f64, w_max: f64, n: usize) -> NoiseGrid {
        let wmin = (w_min * ONE as f64) as i64;
        let wmax = (w_max * ONE as f64) as i64;
        let dw = (wmax - wmin) / (n as i64 - 1).max(1);
        let omega = (0..n).map(|i| wmin + i as i64 * dw).collect();
        NoiseGrid { omega, dw }
    }
}

/// A `1/f^α`-plus-Lorentzian dephasing PSD on the grid (fixed-point). `a_over_f`
/// scales the `1/f` floor; a Lorentzian bump of height `bump` sits at `w0` with
/// half-width `gamma`. This is the hidden spectrum the spectroscopy recovers.
pub fn psd_model(g: &NoiseGrid, a_over_f: f64, bump: f64, w0: f64, gamma: f64) -> Vec<i64> {
    let a = (a_over_f * ONE as f64) as i64;
    let h = (bump * ONE as f64) as i64;
    let w0f = (w0 * ONE as f64) as i64;
    let gf = (gamma * ONE as f64) as i64;
    let g2 = fmul(gf, gf);
    g.omega
        .iter()
        .map(|&w| {
            let one_over_f = fdiv(a, w); // A / ω
            let dwn = w - w0f;
            let lor = fmul(h, fdiv(g2, fmul(dwn, dwn) + g2)); // bump · γ²/((ω−ω0)²+γ²)
            one_over_f + lor
        })
        .collect()
}

fn hash_psd(psd: &[i64]) -> [u8; 32] {
    let mut b = Vec::with_capacity(psd.len() * 8 + 8);
    b.extend_from_slice(b"wai:qc-psd\x01");
    for &v in psd {
        b.extend_from_slice(&v.to_le_bytes());
    }
    content_hash(&b)
}

// ===========================================================================
// DD sequences + filter functions
// ===========================================================================

/// Free induction (Ramsey) — no π-pulses.
pub fn dd_free() -> Vec<i64> {
    Vec::new()
}
/// Hahn echo — one π-pulse at T/2.
pub fn dd_hahn() -> Vec<i64> {
    vec![ONE / 2]
}
/// CPMG-`n` — `n` equally spaced π-pulses at `(k−½)/n`.
pub fn dd_cpmg(n: usize) -> Vec<i64> {
    (1..=n)
        .map(|k| ((2 * k as i64 - 1) * ONE) / (2 * n as i64))
        .collect()
}
/// UDD-`n` — Uhrig spacing `sin²(πk/(2n+2))`, optimal at nulling low-frequency
/// noise to high order.
pub fn dd_udd(n: usize) -> Vec<i64> {
    (1..=n)
        .map(|k| {
            let theta = (PI_FX * k as i64) / (2 * n as i64 + 2);
            let s = sin_fx(theta);
            fmul(s, s)
        })
        .collect()
}

/// The filter kernel of a DD sequence on the grid: `kern[i] = (1/π)·|ỹ(ω_i)|²·dω`,
/// so the coherence overlap against a PSD is simply `χ = Σ_i psd[i]·kern[i]`.
///
/// `ỹ(ω) = ∫₀ᵀ y(t) e^{iωt} dt` with the toggling-frame switching function
/// `y(t) = (−1)^k` between pulses; `|ỹ|² = |Σ_k (−1)^k (e^{iωt_{k+1}} − e^{iωt_k})|² / ω²`.
pub fn filter_kernel(pulses: &[i64], g: &NoiseGrid) -> Vec<i64> {
    // boundary-augmented times: 0, t_1..t_n, T
    let mut t = Vec::with_capacity(pulses.len() + 2);
    t.push(0i64);
    t.extend_from_slice(pulses);
    t.push(ONE);
    g.omega
        .iter()
        .map(|&w| {
            let (mut re, mut im) = (0i64, 0i64);
            // Σ_{k=0}^{n} (-1)^k (E_{k+1} - E_k)
            for k in 0..t.len() - 1 {
                let sign = if k % 2 == 0 { 1i64 } else { -1 };
                let p1 = fmul(w, t[k + 1]);
                let p0 = fmul(w, t[k]);
                re += sign * (cos_fx(p1) - cos_fx(p0));
                im += sign * (sin_fx(p1) - sin_fx(p0));
            }
            // |ỹ|² = (re²+im²)/ω²  (fixed-point, ·ONE); then ·dω·(1/π)
            let mag2 = re as i128 * re as i128 + im as i128 * im as i128;
            let w2 = w as i128 * w as i128;
            let yy = if w2 == 0 { 0 } else { ((mag2 * ONE as i128) / w2) as i64 };
            fmul(fmul(yy, g.dw), INV_PI_FX)
        })
        .collect()
}

/// Coherence-decay overlap `χ = Σ_i psd[i]·kern[i]` (fixed-point).
pub fn overlap(psd: &[i64], kern: &[i64]) -> i64 {
    psd.iter().zip(kern).map(|(&s, &k)| fmul(s, k)).sum()
}

/// Retained coherence `C = exp(−χ)`.
pub fn coherence(chi: i64) -> i64 {
    exp_neg_fx(chi)
}

// ===========================================================================
// Noise spectroscopy — reconstruct S(ω) from a DD filter bank
// ===========================================================================

/// The reconstructed spectrum plus the (hidden) truth it was recovered from.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SpectroResult {
    pub omega: Vec<i64>,
    pub psd_true: Vec<i64>,
    pub psd_est: Vec<i64>,
    /// Number of CPMG sequences in the probe bank.
    pub n_sequences: usize,
    /// Mean absolute reconstruction error over the grid (fixed-point).
    pub error_fx: i64,
}

/// Reconstruct the PSD from a bank of CPMG filters. Each sequence `n = 1..=n_max`
/// is a narrow-band filter; we "measure" its retained coherence `C_n` under the
/// hidden `psd_true` (deterministic, with `shots` of sampling noise), take
/// `χ_n = −ln C_n`, and run `iters` Richardson–Lucy deconvolution steps (which
/// preserve positivity) to recover `S(ω)`.
pub fn reconstruct_psd(
    g: &NoiseGrid,
    psd_true: &[i64],
    n_max: usize,
    shots: u32,
    iters: u32,
    seed: u64,
) -> SpectroResult {
    // Build the filter bank and the "measured" decay exponents.
    let kernels: Vec<Vec<i64>> = (1..=n_max).map(|n| filter_kernel(&dd_cpmg(n), g)).collect();
    let mut chi_meas = Vec::with_capacity(n_max);
    for (idx, kern) in kernels.iter().enumerate() {
        let chi_true = overlap(psd_true, kern);
        let c_true = coherence(chi_true);
        // sample P(+1) = (1+C)/2 with `shots`, estimate Ĉ, then χ̂ = −ln Ĉ
        let mut st = seed.wrapping_mul(0x100_0001).wrapping_add(idx as u64 + 1);
        let p_plus = (ONE + c_true) / 2;
        let mut hits = 0u64;
        for _ in 0..shots {
            let r = (splitmix64(&mut st) >> (64 - FRAC)) as i64;
            if r < p_plus {
                hits += 1;
            }
        }
        let c_hat = (2 * (hits as i128 * ONE as i128) / shots.max(1) as i128 - ONE as i128) as i64;
        let c_hat = c_hat.clamp(ONE / 100, ONE);
        chi_meas.push(-ln_fx(c_hat));
    }

    // Richardson–Lucy: S_i ← S_i · [Σ_n K_{n,i}·(χ_n / (K·S)_n)] / [Σ_n K_{n,i}]
    let ng = g.omega.len();
    let denom: Vec<i64> = (0..ng).map(|i| kernels.iter().map(|k| k[i]).sum()).collect();
    let mut est = vec![ONE / 10; ng]; // flat positive start
    for _ in 0..iters {
        // predicted χ per sequence with current estimate
        let pred: Vec<i64> = kernels.iter().map(|k| overlap(&est, k)).collect();
        let ratio: Vec<i64> = pred
            .iter()
            .zip(&chi_meas)
            .map(|(&p, &m)| if p <= 0 { ONE } else { fdiv(m, p) })
            .collect();
        for i in 0..ng {
            if denom[i] <= 0 {
                continue;
            }
            let num: i64 = kernels.iter().zip(&ratio).map(|(k, &r)| fmul(k[i], r)).sum();
            let corr = fdiv(num, denom[i]);
            est[i] = fmul(est[i], corr).max(0);
        }
    }

    let error_fx = {
        let s: i64 = psd_true
            .iter()
            .zip(&est)
            .map(|(&a, &b)| (a - b).abs())
            .sum();
        s / ng.max(1) as i64
    };

    SpectroResult {
        omega: g.omega.clone(),
        psd_true: psd_true.to_vec(),
        psd_est: est,
        n_sequences: n_max,
        error_fx,
    }
}

impl SpectroResult {
    pub fn hash(&self) -> [u8; 32] {
        let mut h = blake3::Hasher::new();
        h.update(b"wai:qc-spectro\x01");
        for &v in &self.psd_est {
            h.update(&v.to_le_bytes());
        }
        *h.finalize().as_bytes()
    }
    pub fn artifacts(&self) -> CalArtifacts {
        let mut cfg = Vec::new();
        for &v in &self.psd_est {
            cfg.extend_from_slice(&v.to_le_bytes());
        }
        let mut ev = Vec::new();
        for &v in &self.psd_true {
            ev.extend_from_slice(&v.to_le_bytes());
        }
        CalArtifacts {
            target: "noise-psd:q0".into(),
            config_bytes: cfg,
            evidence_kind: "dd_noise_spectroscopy".into(),
            evidence_bytes: ev,
            summary: format!("PSD reconstructed from {} CPMG filters", self.n_sequences),
        }
    }
    pub fn seal(&self, signer: &SigningKey, signer_id: &str, joules_micro: u64, grant: GrantRef) -> CalibrationReceipt {
        seal_artifacts(signer, signer_id, "sim:transmon:q0", &self.artifacts(), joules_micro, grant, None)
    }
}

// ===========================================================================
// Filter-function robust control — reshape a DD sequence to dodge the PSD
// ===========================================================================

/// A robust DD sequence optimised against a target PSD.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RobustResult {
    /// Optimised pulse timing (fixed-point, in `(0,1)`).
    pub pulses: Vec<i64>,
    /// Filter overlap χ of the optimised sequence (lower is better).
    pub chi_opt: i64,
    /// Filter overlap χ of equal-spaced CPMG (the baseline).
    pub chi_cpmg: i64,
    /// Coherence C = exp(−χ) of optimised vs CPMG.
    pub coh_opt: i64,
    pub coh_cpmg: i64,
    /// Hash of the PSD the pulse is robust against (bound into the receipt).
    pub psd_hash: [u8; 32],
    /// Filter kernels for display (optimised, CPMG).
    pub kern_opt: Vec<i64>,
    pub kern_cpmg: Vec<i64>,
}

/// Reshape an `n`-pulse DD sequence to minimise the filter overlap `χ` against
/// `psd`, by deterministic coordinate descent on the pulse positions (starting
/// from CPMG, each position nudged within its neighbours by a shrinking step).
pub fn optimize_dd(g: &NoiseGrid, psd: &[i64], n: usize, rounds: u32) -> RobustResult {
    let cpmg = dd_cpmg(n);
    let kern_cpmg = filter_kernel(&cpmg, g);
    let chi_cpmg = overlap(psd, &kern_cpmg);

    let mut pos = cpmg.clone();
    let chi_of = |p: &[i64]| overlap(psd, &filter_kernel(p, g));
    let mut best = chi_of(&pos);
    let mut step = ONE / 8;
    for _ in 0..rounds {
        for i in 0..n {
            let lo = if i == 0 { ONE / 200 } else { pos[i - 1] + ONE / 200 };
            let hi = if i == n - 1 { ONE - ONE / 200 } else { pos[i + 1] - ONE / 200 };
            for &dir in &[-1i64, 1] {
                let cand = (pos[i] + dir * step).clamp(lo, hi);
                if cand == pos[i] {
                    continue;
                }
                let save = pos[i];
                pos[i] = cand;
                let c = chi_of(&pos);
                if c < best {
                    best = c;
                } else {
                    pos[i] = save;
                }
            }
        }
        step = (step * 3) / 5; // shrink
        if step < ONE / 2000 {
            break;
        }
    }

    let kern_opt = filter_kernel(&pos, g);
    let chi_opt = overlap(psd, &kern_opt);
    RobustResult {
        pulses: pos,
        chi_opt,
        chi_cpmg,
        coh_opt: coherence(chi_opt),
        coh_cpmg: coherence(chi_cpmg),
        psd_hash: hash_psd(psd),
        kern_opt,
        kern_cpmg,
    }
}

impl RobustResult {
    /// Improvement factor χ_cpmg / χ_opt (fixed-point).
    pub fn improvement_fx(&self) -> i64 {
        fdiv(self.chi_cpmg, self.chi_opt.max(1))
    }
    pub fn hash(&self) -> [u8; 32] {
        let mut h = blake3::Hasher::new();
        h.update(b"wai:qc-robust-ff\x01");
        for &p in &self.pulses {
            h.update(&p.to_le_bytes());
        }
        h.update(&self.chi_opt.to_le_bytes());
        *h.finalize().as_bytes()
    }
    pub fn artifacts(&self) -> CalArtifacts {
        let mut cfg = Vec::new();
        for &p in &self.pulses {
            cfg.extend_from_slice(&p.to_le_bytes());
        }
        // evidence binds the PSD hash + achieved overlap
        let mut ev = Vec::new();
        ev.extend_from_slice(&self.psd_hash);
        ev.extend_from_slice(&self.chi_opt.to_le_bytes());
        CalArtifacts {
            target: "robust-dd:q0".into(),
            config_bytes: cfg,
            evidence_kind: "filter_function_robust_control".into(),
            evidence_bytes: ev,
            summary: format!(
                "{}-pulse DD, χ {}→{} vs CPMG (filter dodges the pinned PSD)",
                self.pulses.len(),
                self.chi_cpmg,
                self.chi_opt
            ),
        }
    }
    pub fn seal(&self, signer: &SigningKey, signer_id: &str, joules_micro: u64, grant: GrantRef) -> CalibrationReceipt {
        seal_artifacts(signer, signer_id, "sim:transmon:q0", &self.artifacts(), joules_micro, grant, None)
    }
}

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

    fn key(s: u8) -> SigningKey {
        SigningKey::from_bytes(&[s; 32])
    }
    fn grid() -> NoiseGrid {
        NoiseGrid::linear(0.5, 90.0, 180)
    }
    fn f(v: i64) -> f64 {
        v as f64 / ONE as f64
    }

    #[test]
    fn ln_and_coherence_roundtrip() {
        // exp(-χ) then -ln should recover χ within tolerance.
        for &chi in &[ONE / 10, ONE / 2, ONE, 2 * ONE] {
            let c = coherence(chi);
            let back = -ln_fx(c);
            assert!((back - chi).abs() < ONE / 20, "χ {} → C {} → {}", f(chi), f(c), f(back));
        }
    }

    #[test]
    fn more_pulses_suppress_low_freq_noise() {
        // Against a 1/f spectrum, an echo must retain more coherence than free
        // induction, and CPMG-8 more than a single echo. The physics sanity check.
        let g = grid();
        let psd = psd_model(&g, 0.02, 0.0, 0.0, 1.0);
        let chi_free = overlap(&psd, &filter_kernel(&dd_free(), &g));
        let chi_hahn = overlap(&psd, &filter_kernel(&dd_hahn(), &g));
        let chi_cpmg8 = overlap(&psd, &filter_kernel(&dd_cpmg(8), &g));
        assert!(chi_hahn < chi_free, "echo beats free: {} vs {}", chi_hahn, chi_free);
        assert!(chi_cpmg8 < chi_hahn, "CPMG-8 beats echo: {} vs {}", chi_cpmg8, chi_hahn);
    }

    #[test]
    fn spectroscopy_recovers_hidden_bump() {
        // A Lorentzian bump on a 1/f floor: the reconstruction should place power
        // near the true bump location (correlation of shapes is positive).
        let g = grid();
        let psd = psd_model(&g, 0.015, 0.5, 40.0, 6.0);
        let r = reconstruct_psd(&g, &psd, 20, 40_000, 40, 0xC0FFEE);
        // reconstruction error is a small fraction of the PSD scale
        let mean_psd: i64 = psd.iter().sum::<i64>() / psd.len() as i64;
        assert!(
            r.error_fx < mean_psd,
            "reconstruction error {} should be below mean PSD {}",
            r.error_fx, mean_psd
        );
        // the estimate is positive everywhere (RL preserves positivity)
        assert!(r.psd_est.iter().all(|&v| v >= 0));
    }

    #[test]
    fn spectroscopy_is_deterministic() {
        let g = grid();
        let psd = psd_model(&g, 0.02, 0.3, 30.0, 5.0);
        let a = reconstruct_psd(&g, &psd, 16, 20_000, 30, 7);
        let b = reconstruct_psd(&g, &psd, 16, 20_000, 30, 7);
        assert_eq!(a, b);
    }

    #[test]
    fn robust_control_beats_cpmg() {
        // Reshaping the DD timing must not do worse than CPMG, and against a
        // low-frequency-weighted spectrum it should do strictly better.
        let g = grid();
        let psd = psd_model(&g, 0.04, 0.0, 0.0, 1.0);
        let r = optimize_dd(&g, &psd, 6, 30);
        assert!(r.chi_opt <= r.chi_cpmg, "optimised {} must not exceed CPMG {}", r.chi_opt, r.chi_cpmg);
        assert!(r.coh_opt >= r.coh_cpmg);
    }

    #[test]
    fn results_seal_and_verify() {
        let g = grid();
        let psd = psd_model(&g, 0.02, 0.4, 35.0, 5.0);
        let spec = reconstruct_psd(&g, &psd, 16, 20_000, 30, 1);
        let sr = spec.seal(&key(1), "did:key:lab", 1_000_000, GrantRef::unbounded("quantum.calibrate"));
        assert!(sr.verify());

        let rob = optimize_dd(&g, &psd, 6, 25);
        let rr = rob.seal(&key(2), "did:key:lab", 1_000_000, GrantRef::unbounded("quantum.calibrate"));
        assert!(rr.verify());
        // the robust receipt binds the exact PSD it was optimised against
        assert!(rob.artifacts().evidence_bytes.starts_with(&hash_psd(&psd)));
    }
}