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

//! Matrix Product State backend — the classical referee (`wai.quantum.mps`).
//!
//! Published classical simulations of "utility-scale" circuits (2024–2026) use
//! **tensor networks**: for a state whose entanglement stays bounded, an MPS represents `2ⁿ` amplitudes
//! in `O(n · χ²)` numbers and evolves them gate-by-gate, truncating the bond
//! dimension `χ` at each two-qubit gate. This is the backend the dense simulator's
//! cheapest-sufficient cascade hands large, low-entanglement circuits to.
//!
//! Contract, stated honestly and differently from the dense simulator: an MPS with
//! bounded `χ` is **approximate** — the whole point is to discard small Schmidt
//! weight — so byte-exact-vs-analytic is not the promise here. What IS preserved is
//! **determinism**: the contraction and truncation use only IEEE-754 `+ − × ÷ √`
//! (no FMA), and the phase gate's `cos` and `sin` are computed in-crate rather than
//! by the platform's math library, so the result is bit-identical on every machine
//! and target. The engine reports the **max bond dimension** it used and the **retained weight**
//! (the standard truncation-fidelity heuristic). Verified against the byte-exact dense
//! simulator: exact agreement when `χ` is unbounded, graceful, measured degradation
//! when it is squeezed.

// Numerical kernels: indexed loops and inherent add/mul/sub on the complex type are
// deliberate for legibility of the matrix math.
#![allow(clippy::should_implement_trait, clippy::needless_range_loop, clippy::manual_range_contains)]

use crate::quantum::{BaseGate, Circuit, Gate, StateVector};

pub use crate::linalg::C;
use crate::linalg::jacobi_svd;

// ===========================================================================
// MPS
// ===========================================================================

/// A matrix-product-state on `n` qubits. Site `i`'s tensor `a[i]` has shape
/// `(dl[i], 2, dr[i])`, stored flat with index `(l*2 + s)*dr + r`.
pub struct Mps {
    pub n: u8,
    a: Vec<Vec<C>>,
    dl: Vec<usize>,
    dr: Vec<usize>,
    chi_max: usize,
    max_bond: usize,
    retained: f64, // running product of retained weight (fidelity floor)
}

#[inline]
fn idx(l: usize, s: usize, r: usize, dr: usize) -> usize {
    (l * 2 + s) * dr + r
}

impl Mps {
    /// The product state `|0…0⟩` on `n` qubits.
    pub fn zero(n: u8, chi_max: usize) -> Mps {
        let a: Vec<Vec<C>> = (0..n).map(|_| vec![C::ONE, C::ZERO]).collect(); // (1,2,1): s=0→1
        Mps {
            n,
            a,
            dl: vec![1; n as usize],
            dr: vec![1; n as usize],
            chi_max: chi_max.max(1),
            max_bond: 1,
            retained: 1.0,
        }
    }

    /// Apply a single-qubit gate (2×2, `g[out][in]`) on site `i`.
    fn apply1(&mut self, i: usize, g: [[C; 2]; 2]) {
        let (dl, dr) = (self.dl[i], self.dr[i]);
        let ai = &mut self.a[i];
        for l in 0..dl {
            for r in 0..dr {
                let a0 = ai[idx(l, 0, r, dr)];
                let a1 = ai[idx(l, 1, r, dr)];
                ai[idx(l, 0, r, dr)] = g[0][0].mul(a0).add(g[0][1].mul(a1));
                ai[idx(l, 1, r, dr)] = g[1][0].mul(a0).add(g[1][1].mul(a1));
            }
        }
    }

    /// Apply a two-qubit gate (4×4, basis `2*s_lo + s_hi`) on **adjacent** sites
    /// `i, i+1`, then SVD-truncate the bond to `chi_max`.
    fn apply2_adjacent(&mut self, i: usize, g: [[C; 4]; 4]) {
        let (dl, dc, dr) = (self.dl[i], self.dr[i], self.dr[i + 1]);
        debug_assert_eq!(dc, self.dl[i + 1]);
        // contract A[i]·A[i+1] over the shared bond → Theta[l, s1, s2, r]
        // apply gate: Theta'[l,t1,t2,r] = sum G[(t1,t2)][(s1,s2)] Theta[l,s1,s2,r]
        // reshape to M[row=(l,t1), col=(t2,r)]  (rows = dl*2, cols = 2*dr)
        let rows = dl * 2;
        let cols = 2 * dr;
        let mut mcols: Vec<Vec<C>> = vec![vec![C::ZERO; rows]; cols];
        for l in 0..dl {
            for r in 0..dr {
                // gather Theta[l, ., ., r] as a length-4 vector over (s1,s2)
                let mut th = [C::ZERO; 4];
                for s1 in 0..2 {
                    for s2 in 0..2 {
                        let mut acc = C::ZERO;
                        for c in 0..dc {
                            acc = acc.add(self.a[i][idx(l, s1, c, dc)].mul(self.a[i + 1][idx(c, s2, r, dr)]));
                        }
                        th[s1 * 2 + s2] = acc;
                    }
                }
                for t1 in 0..2 {
                    for t2 in 0..2 {
                        let mut val = C::ZERO;
                        for s in 0..4 {
                            val = val.add(g[t1 * 2 + t2][s].mul(th[s]));
                        }
                        let row = l * 2 + t1;
                        let col = t2 * dr + r;
                        mcols[col][row] = val;
                    }
                }
            }
        }
        let (u, s, v) = jacobi_svd(mcols, rows, cols);
        // truncate
        let total: f64 = s.iter().map(|x| x * x).sum();
        let mut chi = 0;
        for (k, sv) in s.iter().enumerate() {
            if k >= self.chi_max || *sv <= 1e-12 * s[0].max(1e-300) {
                break;
            }
            chi += 1;
        }
        let chi = chi.max(1).min(s.len());
        let kept: f64 = s.iter().take(chi).map(|x| x * x).sum();
        if total > 0.0 {
            self.retained *= kept / total;
        }
        self.max_bond = self.max_bond.max(chi);
        // new A[i] = U[:, :chi]  → shape (dl, 2, chi)
        let mut na = vec![C::ZERO; dl * 2 * chi];
        for l in 0..dl {
            for t1 in 0..2 {
                for k in 0..chi {
                    na[idx(l, t1, k, chi)] = u[k][l * 2 + t1];
                }
            }
        }
        // new A[i+1] = diag(S)·V†  → shape (chi, 2, dr): [k,t2,r] = S[k]·conj(V_col[k][t2*dr+r])
        let mut nb = vec![C::ZERO; chi * 2 * dr];
        for k in 0..chi {
            for t2 in 0..2 {
                for r in 0..dr {
                    nb[idx(k, t2, r, dr)] = v[k][t2 * dr + r].conj().scale(s[k]);
                }
            }
        }
        self.a[i] = na;
        self.dr[i] = chi;
        self.a[i + 1] = nb;
        self.dl[i + 1] = chi;
    }

    /// Adjacent SWAP of sites `i, i+1` (its own two-qubit gate).
    fn swap_adjacent(&mut self, i: usize) {
        self.apply2_adjacent(i, swap_gate());
    }

    /// Apply a two-qubit gate on sites `a`,`b` (any distance) by SWAP-routing `b`
    /// next to `a`, applying, and SWAPping back. `g4` is in `(min,max)`-site order.
    fn apply2(&mut self, a: usize, b: usize, g4: [[C; 4]; 4]) {
        let (lo, hi) = (a.min(b), a.max(b));
        if hi == lo + 1 {
            self.apply2_adjacent(lo, g4);
            return;
        }
        // bring site `hi` down to `lo+1`
        for k in (lo + 1..hi).rev() {
            self.swap_adjacent(k);
        }
        self.apply2_adjacent(lo, g4);
        for k in lo + 1..hi {
            self.swap_adjacent(k);
        }
    }

    /// The full `2ⁿ` statevector (for verification / small `n`), in the dense
    /// simulator's index convention (bit `i` = qubit/site `i`), **normalized**.
    pub fn to_statevector(&self) -> Vec<C> {
        let dim = 1usize << self.n;
        let mut out = vec![C::ZERO; dim];
        for k in 0..dim {
            // contract left→right, selecting s_i = (k>>i)&1
            let mut vec = vec![C::ONE]; // 1×1 row
            for i in 0..self.n as usize {
                let s = (k >> i) & 1;
                let (dl, dr) = (self.dl[i], self.dr[i]);
                let mut nv = vec![C::ZERO; dr];
                for r in 0..dr {
                    let mut acc = C::ZERO;
                    for l in 0..dl {
                        acc = acc.add(vec[l].mul(self.a[i][idx(l, s, r, dr)]));
                    }
                    nv[r] = acc;
                }
                vec = nv;
            }
            out[k] = vec[0];
        }
        // normalize
        let norm: f64 = out.iter().map(|c| c.norm2()).sum::<f64>().sqrt();
        if norm > 0.0 {
            for c in out.iter_mut() {
                *c = c.scale(1.0 / norm);
            }
        }
        out
    }

    pub fn max_bond(&self) -> usize {
        self.max_bond
    }
    /// Retained Schmidt weight across all truncations, in `[0,1]` — the standard
    /// truncation-fidelity heuristic (the fraction of squared singular weight kept).
    /// It is an *indicator*, not a strict bound: because discarded components can be
    /// amplified by later gates, the true fidelity can sit either side of it. `1.0`
    /// means no weight was ever discarded (the run was exact).
    pub fn retained_weight(&self) -> f64 {
        self.retained
    }
}

// ===========================================================================
// Gate matrices (f64) + circuit runner
// ===========================================================================

fn m1(base: BaseGate, param: u16) -> [[C; 2]; 2] {
    let s = std::f64::consts::FRAC_1_SQRT_2;
    let z = C::ZERO;
    let o = C::ONE;
    let i = C::new(0.0, 1.0);
    let ni = C::new(0.0, -1.0);
    let no = C::new(-1.0, 0.0);
    let phase = |k: u16| {
        let th = 2.0 * std::f64::consts::PI / (1u64 << k) as f64;
        let (s, c) = crate::repro::sin_cos(th);
        C::new(c, s)
    };
    match base {
        BaseGate::I => [[o, z], [z, o]],
        BaseGate::X => [[z, o], [o, z]],
        BaseGate::Y => [[z, ni], [i, z]],
        BaseGate::Z => [[o, z], [z, no]],
        BaseGate::H => [[C::new(s, 0.0), C::new(s, 0.0)], [C::new(s, 0.0), C::new(-s, 0.0)]],
        BaseGate::S => [[o, z], [z, i]],
        BaseGate::Sdg => [[o, z], [z, ni]],
        BaseGate::T => [[o, z], [z, phase(3)]],
        BaseGate::Tdg => [[o, z], [z, phase(3).conj()]],
        BaseGate::P => [[o, z], [z, phase(param)]],
    }
}

/// The 4×4 for a controlled-`U` on sites `(lo,hi)` (basis `2*s_lo+s_hi`), given
/// which of the pair is the control.
fn controlled4(control_is_lo: bool, u: [[C; 2]; 2]) -> [[C; 4]; 4] {
    let mut g = [[C::ZERO; 4]; 4];
    for a in 0..2 {
        for b in 0..2 {
            let inb = a * 2 + b;
            let (ctrl, tgt) = if control_is_lo { (a, b) } else { (b, a) };
            if ctrl == 0 {
                g[inb][inb] = C::ONE; // identity
            } else {
                // apply U to the target bit
                for tt in 0..2 {
                    let (oa, ob) = if control_is_lo { (a, tt) } else { (tt, b) };
                    let out = oa * 2 + ob;
                    g[out][inb] = u[tt][tgt];
                }
            }
        }
    }
    g
}

fn swap_gate() -> [[C; 4]; 4] {
    let mut g = [[C::ZERO; 4]; 4];
    g[0][0] = C::ONE;
    g[3][3] = C::ONE;
    g[1][2] = C::ONE; // |01> <-> |10>
    g[2][1] = C::ONE;
    g
}

/// Expand a Toffoli-class gate (`X`/`Z` with two controls) into 1- and 2-qubit
/// gates the MPS handles directly (standard Clifford+T Toffoli; CCZ = H·CCX·H).
fn decompose_cc(g: &Gate) -> Option<Vec<Gate>> {
    if g.controls.len() != 2 {
        return None;
    }
    let (a, b, c) = (g.controls[0], g.controls[1], g.target);
    let g1 = |base, t| Gate { base, controls: vec![], target: t, param: 0 };
    let cx = |ctrl, t| Gate { base: BaseGate::X, controls: vec![ctrl], target: t, param: 0 };
    let mut ops = Vec::new();
    if g.base == BaseGate::Z {
        ops.push(g1(BaseGate::H, c));
    }
    // CCX on (a,b;c)
    ops.push(g1(BaseGate::H, c));
    ops.push(cx(b, c));
    ops.push(g1(BaseGate::Tdg, c));
    ops.push(cx(a, c));
    ops.push(g1(BaseGate::T, c));
    ops.push(cx(b, c));
    ops.push(g1(BaseGate::Tdg, c));
    ops.push(cx(a, c));
    ops.push(g1(BaseGate::T, b));
    ops.push(g1(BaseGate::T, c));
    ops.push(cx(a, b));
    ops.push(g1(BaseGate::H, c));
    ops.push(g1(BaseGate::T, a));
    ops.push(g1(BaseGate::Tdg, b));
    ops.push(cx(a, b));
    if g.base == BaseGate::Z {
        ops.push(g1(BaseGate::H, c));
    }
    Some(ops)
}

/// The result of an MPS run.
pub struct MpsRun {
    pub max_bond: usize,
    pub retained_weight: f64,
    /// Fidelity `|⟨mps|dense⟩|²` against the byte-exact dense statevector, when the
    /// caller asks to verify (small `n`); `None` if not computed.
    pub fidelity_vs_dense: Option<f64>,
}

/// Errors an MPS run can raise.
#[derive(Debug)]
pub enum MpsError {
    /// A gate with more than two controls that isn't a decomposable Toffoli class.
    Unsupported(String),
}

/// Run `circuit` as an MPS with bond cap `chi_max`. Returns the final [`Mps`] and a
/// stats summary. Gates with >2 controls (other than CCX/CCZ) are unsupported.
pub fn run(circuit: &Circuit, chi_max: usize) -> Result<(Mps, MpsRun), MpsError> {
    let mut mps = Mps::zero(circuit.n_qubits, chi_max);
    for g in &circuit.ops {
        apply_gate(&mut mps, g)?;
    }
    let run = MpsRun {
        max_bond: mps.max_bond(),
        retained_weight: mps.retained_weight(),
        fidelity_vs_dense: None,
    };
    Ok((mps, run))
}

fn apply_gate(mps: &mut Mps, g: &Gate) -> Result<(), MpsError> {
    match g.controls.len() {
        0 => mps.apply1(g.target as usize, m1(g.base, g.param)),
        1 => {
            let ctrl = g.controls[0] as usize;
            let tgt = g.target as usize;
            let control_is_lo = ctrl < tgt;
            let g4 = controlled4(control_is_lo, m1(g.base, g.param));
            mps.apply2(ctrl, tgt, g4);
        }
        2 => {
            let ops = decompose_cc(g)
                .ok_or_else(|| MpsError::Unsupported(format!("{:?} with 2 controls", g.base)))?;
            for sub in &ops {
                apply_gate(mps, sub)?;
            }
        }
        k => return Err(MpsError::Unsupported(format!("{k} controls"))),
    }
    Ok(())
}

/// Fidelity `|⟨mps|dense⟩|²` between an MPS statevector and the byte-exact dense
/// statevector — the verification metric.
pub fn fidelity(mps: &Mps, dense: &StateVector) -> f64 {
    let a = mps.to_statevector();
    let scale = (1u64 << crate::quantum::FRAC) as f64;
    let mut overlap = C::ZERO;
    let mut nd = 0.0;
    for (k, amp) in dense.amps.iter().enumerate() {
        let d = C::new(amp.re as f64 / scale, amp.im as f64 / scale);
        overlap = overlap.add(a[k].conj().mul(d));
        nd += d.norm2();
    }
    if nd <= 0.0 {
        return 0.0;
    }
    overlap.norm2() / nd // |<mps|dense>|^2 (mps already normalized)
}

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

    fn approx(a: f64, b: f64, tol: f64) -> bool {
        (a - b).abs() < tol
    }
    fn ghz(n: u8) -> Circuit {
        let mut c = Circuit::new(n);
        c.h(0);
        for i in 1..n {
            c.cx(i - 1, i);
        }
        c
    }
    fn bell() -> Circuit {
        let mut c = Circuit::new(2);
        c.h(0);
        c.cx(0, 1);
        c
    }
    /// A brickwork circuit that genuinely builds chain entanglement (bond → 2^(n/2)).
    fn entangling(n: u8) -> Circuit {
        let mut c = Circuit::new(n);
        for _ in 0..3 {
            for q in 0..n {
                c.h(q);
                c.t(q);
            }
            let mut q = 0;
            while q + 1 < n {
                c.cx(q, q + 1);
                q += 2;
            }
            let mut q = 1;
            while q + 1 < n {
                c.cx(q, q + 1);
                q += 2;
            }
        }
        c
    }

    #[test]
    fn product_state_stays_bond_one() {
        // all-H on |0…0> is a product state: MPS bond never exceeds 1.
        let mut c = Circuit::new(5);
        for q in 0..5 {
            c.h(q);
        }
        let (mps, run) = run(&c, 16).unwrap();
        assert_eq!(run.max_bond, 1, "product state must stay bond-1");
        assert!(approx(fidelity(&mps, &c.simulate().unwrap()), 1.0, 1e-9));
    }

    #[test]
    fn ghz_is_exact_at_bond_two() {
        let c = ghz(6);
        let (mps, run) = run(&c, 8).unwrap();
        assert_eq!(run.max_bond, 2, "GHZ is a bond-2 state");
        assert!(approx(fidelity(&mps, &c.simulate().unwrap()), 1.0, 1e-9), "GHZ must be exact");
    }

    #[test]
    fn bell_and_small_circuits_match_dense() {
        for c in [bell(), ghz(4)] {
            let (mps, _) = run(&c, 16).unwrap();
            let f = fidelity(&mps, &c.simulate().unwrap());
            assert!(approx(f, 1.0, 1e-8), "MPS should match dense, got {f}");
        }
    }

    #[test]
    fn qft_matches_dense_at_full_bond() {
        // QFT builds entanglement; at full bond (2^(n/2)) MPS is exact.
        let c = Circuit::qft(6);
        let (mps, _) = run(&c, 64).unwrap();
        let f = fidelity(&mps, &c.simulate().unwrap());
        assert!(approx(f, 1.0, 1e-7), "QFT full-bond fidelity {f}");
    }

    #[test]
    fn ccz_decomposition_matches_dense() {
        // a 2-control CCZ over |+++>: dense applies it natively, MPS decomposes it
        // into 1/2-qubit gates — they must agree.
        let mut c = Circuit::new(3);
        c.h(0);
        c.h(1);
        c.h(2);
        c.ops.push(Gate { base: BaseGate::Z, controls: vec![0, 1], target: 2, param: 0 });
        let (mps, _) = run(&c, 32).unwrap();
        let f = fidelity(&mps, &c.simulate().unwrap());
        assert!(approx(f, 1.0, 1e-6), "CCZ decomposition fidelity {f}");
    }

    #[test]
    fn nonadjacent_two_qubit_gate() {
        // CX between distant sites (0 and 4) must route via SWAPs and still match.
        let mut c = Circuit::new(5);
        c.h(0);
        c.cx(0, 4);
        let (mps, _) = run(&c, 8).unwrap();
        assert!(approx(fidelity(&mps, &c.simulate().unwrap()), 1.0, 1e-8));
    }

    #[test]
    fn truncation_is_measured_not_hidden() {
        // Squeeze QFT's bond below what it needs → the engine REPORTS the lost
        // weight (retained_weight < 1) and the loss is REAL (fidelity < 1). The
        // truncation is neither hidden nor free.
        let c = entangling(6);
        let dense = c.simulate().unwrap();
        let (mps, st) = run(&c, 2).unwrap();
        assert!(st.retained_weight < 0.999, "squeezed bond must report lost weight: {}", st.retained_weight);
        assert!(st.retained_weight > 0.0, "some weight is kept");
        assert!(st.max_bond <= 2, "bond cap respected");
        assert!(fidelity(&mps, &dense) < 0.999, "aggressive truncation must measurably reduce fidelity");
        // and a generous bond (full = 2^(n/2) = 8 for n=6) recovers it exactly
        let (mps2, st2) = run(&c, 8).unwrap();
        assert!(approx(st2.retained_weight, 1.0, 1e-6), "full bond keeps all weight: {}", st2.retained_weight);
        assert!(approx(fidelity(&mps2, &dense), 1.0, 1e-6));
    }

    #[test]
    fn deterministic() {
        let c = Circuit::qft(6);
        let a = run(&c, 8).unwrap().0.to_statevector();
        let b = run(&c, 8).unwrap().0.to_statevector();
        assert!(a.iter().zip(&b).all(|(x, y)| x.re.to_bits() == y.re.to_bits() && x.im.to_bits() == y.im.to_bits()),
            "MPS must be bit-reproducible");
    }
}