prism-q 0.28.0

//! Pauli propagation engines for Clifford+T circuits: SPP samples backward
//! Heisenberg propagation stochastically, SPD carries the full weighted Pauli
//! sum with optional truncation. Neither materializes a state vector.

use num_complex::Complex64;
use rand::SeedableRng;
use rand::{Rng, RngExt};
use rand_chacha::ChaCha8Rng;

use crate::circuit::{Circuit, Instruction, SmallVec};
use crate::error::{PrismError, Result};
use crate::gates::Gate;
use crate::sim::compiled::{PauliVec, flip_bit, propagate_backward};

const SQRT_2: f64 = std::f64::consts::SQRT_2;

/// Absolute ceiling on the weighted-Pauli term count, enforced even in exact
/// mode (`max_terms == 0`). Without a per-step truncation budget the term set
/// grows as `2^(in-cone T count)`; this caps transient memory the same way the
/// stabilizer-rank backend does. Exceeding it is an error, not a silent
/// truncation: callers wanting bounded approximate evaluation pass a nonzero
/// `max_terms` instead.
const SPD_MAX_TERMS_CEILING: usize = 1 << 20;

#[inline]
fn check_spd_term_ceiling(len: usize) -> Result<()> {
    if len > SPD_MAX_TERMS_CEILING {
        return Err(PrismError::BackendUnsupported {
            backend: "SPD".into(),
            operation: format!(
                "weighted-Pauli sum exceeded {SPD_MAX_TERMS_CEILING} terms; pass a nonzero \
                 max_terms for bounded approximate truncation"
            ),
        });
    }
    Ok(())
}

fn validate_clifford_t_unitary(circuit: &Circuit, backend: &'static str) -> Result<()> {
    for inst in &circuit.instructions {
        match inst {
            Instruction::Gate { gate, .. } => {
                if !(gate.is_clifford() || matches!(gate, Gate::T | Gate::Tdg)) {
                    return Err(PrismError::BackendUnsupported {
                        backend: backend.to_string(),
                        operation: format!("non-Clifford+T gate `{}`", gate.name()),
                    });
                }
            }
            Instruction::Barrier { .. } => {}
            Instruction::Measure { .. }
            | Instruction::Reset { .. }
            | Instruction::Conditional { .. } => {
                return Err(PrismError::IncompatibleBackend {
                    backend: backend.to_string(),
                    reason: "Pauli propagation requires a unitary Clifford+T circuit without measurements, resets, or conditionals"
                        .to_string(),
                });
            }
        }
    }
    Ok(())
}

// ---- Coalesced circuit representation ----

enum CoalescedOp {
    SmallCliff(Vec<(Gate, SmallVec<[usize; 4]>)>),
    T { qubit: usize, is_dagger: bool },
}

fn coalesce_cliffords(circuit: &Circuit) -> Vec<CoalescedOp> {
    let mut ops = Vec::new();
    let mut cliff_buf: Vec<(Gate, SmallVec<[usize; 4]>)> = Vec::new();

    for inst in &circuit.instructions {
        if let Instruction::Gate { gate, targets } = inst {
            match gate {
                Gate::T => {
                    flush_cliff_buf(&mut cliff_buf, &mut ops);
                    ops.push(CoalescedOp::T {
                        qubit: targets[0],
                        is_dagger: false,
                    });
                }
                Gate::Tdg => {
                    flush_cliff_buf(&mut cliff_buf, &mut ops);
                    ops.push(CoalescedOp::T {
                        qubit: targets[0],
                        is_dagger: true,
                    });
                }
                _ => {
                    cliff_buf.push((gate.clone(), SmallVec::from_slice(targets)));
                }
            }
        } else {
            flush_cliff_buf(&mut cliff_buf, &mut ops);
        }
    }
    flush_cliff_buf(&mut cliff_buf, &mut ops);
    ops
}

fn flush_cliff_buf(buf: &mut Vec<(Gate, SmallVec<[usize; 4]>)>, ops: &mut Vec<CoalescedOp>) {
    if buf.is_empty() {
        return;
    }
    ops.push(CoalescedOp::SmallCliff(std::mem::take(buf)));
}

// ---- SPP (Stochastic Pauli Propagation) ----

#[inline(always)]
fn branch_t_gate(
    pauli: &mut PauliVec,
    qubit: usize,
    is_dagger: bool,
    rng: &mut impl Rng,
) -> Complex64 {
    // PauliVec stores the Y letter as the ordered product XZ. Since
    // actual Y = i XZ, the T flip branch is imaginary in this basis:
    // T contributes -i/sqrt(2), and Tdg contributes +i/sqrt(2).
    // SPP samples one branch with probability 1/2, so the per-shot
    // flip weight is +/-i*sqrt(2).
    if !pauli.has_x_or_y(qubit) {
        return Complex64::new(1.0, 0.0);
    }

    let keep = rng.random_bool(0.5);
    if keep {
        return Complex64::new(SQRT_2, 0.0);
    }

    flip_bit(&mut pauli.z, qubit);
    if is_dagger {
        Complex64::new(0.0, SQRT_2)
    } else {
        Complex64::new(0.0, -SQRT_2)
    }
}

fn backward_propagate_coalesced(
    ops: &[CoalescedOp],
    observable: &PauliVec,
    rng: &mut impl Rng,
) -> (PauliVec, Complex64) {
    let mut pauli = PauliVec {
        x: observable.x.clone(),
        z: observable.z.clone(),
    };
    let mut weight = Complex64::new(1.0, 0.0);

    for op in ops.iter().rev() {
        match op {
            CoalescedOp::SmallCliff(gates) => {
                for (gate, targets) in gates.iter().rev() {
                    // Phase track Clifford conjugation, mirroring the
                    // SPD path (`conjugate_all_backward_phased`); the
                    // bare `propagate_backward` is Pauli-frame only
                    // and drops the `-1` from e.g. `HYH = -Y`.
                    weight *= clifford_conjugation_phase(gate, targets, &pauli);
                    propagate_backward(&mut pauli, gate, targets);
                }
            }
            CoalescedOp::T { qubit, is_dagger } => {
                weight *= branch_t_gate(&mut pauli, *qubit, *is_dagger, rng);
            }
        }
    }

    (pauli, weight)
}

fn count_t_gates(circuit: &Circuit) -> usize {
    circuit
        .instructions
        .iter()
        .filter(|inst| {
            matches!(
                inst,
                Instruction::Gate {
                    gate: Gate::T | Gate::Tdg,
                    ..
                }
            )
        })
        .count()
}

/// Per-qubit `⟨Z_q⟩` estimates from a stochastic Pauli propagation run.
pub struct SppResult {
    pub expectations: Vec<f64>,
    pub std_errors: Vec<f64>,
    /// Samples drawn per qubit, not in total.
    pub num_samples: usize,
    /// T plus Tdg gates in the circuit.
    pub t_count: usize,
    /// Fraction of samples whose propagated Pauli was diagonal (contributed a
    /// nonzero value).
    pub nonzero_fraction: f64,
}

fn estimate_qubit_expectation(
    ops: &[CoalescedOp],
    qubit: usize,
    num_words: usize,
    num_samples: usize,
    seed: u64,
) -> (f64, f64, usize) {
    let obs = PauliVec::z_on_qubit(num_words, qubit);
    let mut rng = ChaCha8Rng::seed_from_u64(seed);

    let mut mean = 0.0f64;
    let mut m2 = 0.0f64;
    let mut nonzero = 0usize;

    for i in 0..num_samples {
        let (pauli, weight) = backward_propagate_coalesced(ops, &obs, &mut rng);
        let val = if pauli.is_diagonal() {
            nonzero += 1;
            weight.re
        } else {
            0.0
        };

        let delta = val - mean;
        mean += delta / (i + 1) as f64;
        let delta2 = val - mean;
        m2 += delta * delta2;
    }

    let variance = if num_samples > 1 {
        m2 / (num_samples - 1) as f64
    } else {
        0.0
    };
    let std_error = (variance / num_samples as f64).sqrt();
    (mean, std_error, nonzero)
}

/// Pauli axis for a joint-observable term. Identity factors are omitted
/// from the term list and contribute trivially.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum PauliAxis {
    X,
    Y,
    Z,
}

/// One non-identity factor of a joint Pauli observable.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct PauliTerm {
    pub qubit: usize,
    pub axis: PauliAxis,
}

impl PauliTerm {
    pub fn new(qubit: usize, axis: PauliAxis) -> Self {
        Self { qubit, axis }
    }

    pub fn x(qubit: usize) -> Self {
        Self::new(qubit, PauliAxis::X)
    }

    pub fn y(qubit: usize) -> Self {
        Self::new(qubit, PauliAxis::Y)
    }

    pub fn z(qubit: usize) -> Self {
        Self::new(qubit, PauliAxis::Z)
    }
}

/// Result of running SPP on a joint Pauli observable.
#[derive(Debug, Clone)]
pub struct SppObservableResult {
    pub mean: f64,
    pub std_error: f64,
    /// Sample variance of the per-sample contributions, not of the mean.
    pub variance: f64,
    pub num_samples: usize,
    /// Fraction of samples whose propagated Pauli was diagonal.
    pub nonzero_fraction: f64,
    pub t_count: usize,
}

/// Result of running SPD on a joint Pauli observable.
#[derive(Debug, Clone)]
pub struct SpdObservableResult {
    /// `⟨P⟩`; exact when nothing was truncated.
    pub mean: f64,
    pub t_count: usize,
    /// Peak size of the weighted Pauli sum during propagation.
    pub peak_terms: usize,
    /// Total coefficient magnitude discarded by truncation (0 for exact runs).
    pub total_discarded: f64,
}

fn pauli_vec_from_terms(
    num_qubits: usize,
    terms: &[PauliTerm],
) -> std::result::Result<(PauliVec, Complex64), crate::error::PrismError> {
    let num_words = num_qubits.div_ceil(64);
    let mut pv = PauliVec::new(num_words);
    let mut coeff = Complex64::new(1.0, 0.0);
    let mut seen = vec![false; num_qubits];
    for term in terms {
        if term.qubit >= num_qubits {
            return Err(crate::error::PrismError::InvalidQubit {
                index: term.qubit,
                register_size: num_qubits,
            });
        }
        if seen[term.qubit] {
            return Err(crate::error::PrismError::InvalidParameter {
                message: format!(
                    "joint Pauli observable has duplicate factor on qubit {}",
                    term.qubit
                ),
            });
        }
        seen[term.qubit] = true;
        match term.axis {
            PauliAxis::X => {
                pv.x[term.qubit / 64] |= 1u64 << (term.qubit % 64);
            }
            PauliAxis::Z => {
                pv.z[term.qubit / 64] |= 1u64 << (term.qubit % 64);
            }
            PauliAxis::Y => {
                pv.x[term.qubit / 64] |= 1u64 << (term.qubit % 64);
                pv.z[term.qubit / 64] |= 1u64 << (term.qubit % 64);
                coeff *= Complex64::new(0.0, 1.0);
            }
        }
    }
    Ok((pv, coeff))
}

/// Estimate `⟨0^n| U† P U |0^n⟩` for joint Pauli observable `P` on a
/// Clifford+T circuit `U` via stochastic Pauli propagation.
///
/// Each sample backward-propagates the observable through the circuit
/// (Clifford segments as coalesced gate runs, T gates via a stochastic
/// Pauli branch that records a complex weight). The
/// contribution is `Re(weight)` when the final Pauli is diagonal in
/// `{I, Z}` (i.e. evaluates trivially on `|0^n⟩`), else zero.
pub fn run_spp_observable(
    circuit: &Circuit,
    observable: &[PauliTerm],
    num_samples: usize,
    seed: u64,
) -> Result<SppObservableResult> {
    validate_clifford_t_unitary(circuit, "SPP observable")?;
    let n = circuit.num_qubits;
    let t_count = count_t_gates(circuit);
    let ops = coalesce_cliffords(circuit);
    let (obs, obs_coeff) = pauli_vec_from_terms(n, observable)?;

    let mut rng = ChaCha8Rng::seed_from_u64(seed);
    let mut mean = 0.0f64;
    let mut m2 = 0.0f64;
    let mut nonzero = 0usize;

    for i in 0..num_samples {
        let (pauli, weight) = backward_propagate_coalesced(&ops, &obs, &mut rng);
        let val = if pauli.is_diagonal() {
            nonzero += 1;
            (obs_coeff * weight).re
        } else {
            0.0
        };
        let delta = val - mean;
        mean += delta / (i + 1) as f64;
        let delta2 = val - mean;
        m2 += delta * delta2;
    }

    let variance = if num_samples > 1 {
        m2 / (num_samples - 1) as f64
    } else {
        0.0
    };
    let std_error = (variance / num_samples.max(1) as f64).sqrt();
    let nonzero_fraction = nonzero as f64 / num_samples.max(1) as f64;

    Ok(SppObservableResult {
        mean,
        std_error,
        variance,
        num_samples,
        nonzero_fraction,
        t_count,
    })
}

/// Estimate `⟨Z_q⟩` for every qubit of a unitary Clifford+T circuit via
/// stochastic Pauli propagation.
///
/// Draws `num_samples` backward propagations per qubit; the cost scales with
/// samples and gate count, not `2^n`. See [`run_spp_observable`] for a single
/// joint observable.
///
/// # Examples
///
/// ```
/// use prism_q::{Circuit, Gate, run_spp};
///
/// let mut circuit = Circuit::new(2, 0);
/// circuit.add_gate(Gate::H, &[0]);
/// circuit.add_gate(Gate::T, &[0]);
/// circuit.add_gate(Gate::H, &[0]);
///
/// let result = run_spp(&circuit, 20_000, 42)?;
/// assert_eq!(result.t_count, 1);
/// // <Z0> = cos(pi/4); qubit 1 is untouched, so <Z1> = 1 exactly.
/// assert!((result.expectations[0] - std::f64::consts::FRAC_1_SQRT_2).abs() < 0.05);
/// assert!((result.expectations[1] - 1.0).abs() < 1e-12);
/// # Ok::<(), prism_q::PrismError>(())
/// ```
pub fn run_spp(circuit: &Circuit, num_samples: usize, seed: u64) -> Result<SppResult> {
    validate_clifford_t_unitary(circuit, "SPP")?;
    let n = circuit.num_qubits;
    let num_words = n.div_ceil(64);
    let t_count = count_t_gates(circuit);
    let ops = coalesce_cliffords(circuit);

    #[cfg(feature = "parallel")]
    let results: Vec<(f64, f64, usize)> = {
        use rayon::prelude::*;
        (0..n)
            .into_par_iter()
            .map(|q| {
                estimate_qubit_expectation(
                    &ops,
                    q,
                    num_words,
                    num_samples,
                    seed.wrapping_add(q as u64),
                )
            })
            .collect()
    };

    #[cfg(not(feature = "parallel"))]
    let results: Vec<(f64, f64, usize)> = (0..n)
        .map(|q| {
            estimate_qubit_expectation(&ops, q, num_words, num_samples, seed.wrapping_add(q as u64))
        })
        .collect();

    let mut expectations = Vec::with_capacity(n);
    let mut std_errors = Vec::with_capacity(n);
    let mut total_nonzero = 0usize;

    for (mean, std_error, nonzero) in &results {
        expectations.push(*mean);
        std_errors.push(*std_error);
        total_nonzero += nonzero;
    }

    let nonzero_fraction = total_nonzero as f64 / (n * num_samples) as f64;

    Ok(SppResult {
        expectations,
        std_errors,
        num_samples,
        t_count,
        nonzero_fraction,
    })
}

#[cfg(test)]
fn spp_to_probabilities(result: &SppResult) -> Vec<f64> {
    super::expectations_to_marginals(&result.expectations)
        .into_iter()
        .flat_map(|(p0, p1)| [p0, p1])
        .collect()
}

// ---- Deterministic Sparse Pauli Dynamics (SPD) ----

use std::collections::HashMap;

#[inline(always)]
fn pv_get_bit(words: &[u64], qubit: usize) -> bool {
    (words[qubit / 64] >> (qubit % 64)) & 1 != 0
}

#[inline(always)]
fn clifford_conjugation_phase(gate: &Gate, targets: &[usize], pauli: &PauliVec) -> Complex64 {
    match gate {
        Gate::H => {
            let q = targets[0];
            if pv_get_bit(&pauli.x, q) && pv_get_bit(&pauli.z, q) {
                Complex64::new(-1.0, 0.0)
            } else {
                Complex64::new(1.0, 0.0)
            }
        }
        Gate::S => {
            let q = targets[0];
            if pv_get_bit(&pauli.x, q) {
                Complex64::new(0.0, -1.0)
            } else {
                Complex64::new(1.0, 0.0)
            }
        }
        Gate::Sdg => {
            let q = targets[0];
            if pv_get_bit(&pauli.x, q) {
                Complex64::new(0.0, 1.0)
            } else {
                Complex64::new(1.0, 0.0)
            }
        }
        Gate::X => {
            let q = targets[0];
            if pv_get_bit(&pauli.z, q) {
                Complex64::new(-1.0, 0.0)
            } else {
                Complex64::new(1.0, 0.0)
            }
        }
        Gate::Y => {
            let q = targets[0];
            let xq = pv_get_bit(&pauli.x, q);
            let zq = pv_get_bit(&pauli.z, q);
            if xq ^ zq {
                Complex64::new(-1.0, 0.0)
            } else {
                Complex64::new(1.0, 0.0)
            }
        }
        Gate::Z => {
            let q = targets[0];
            if pv_get_bit(&pauli.x, q) {
                Complex64::new(-1.0, 0.0)
            } else {
                Complex64::new(1.0, 0.0)
            }
        }
        Gate::SX => {
            let q = targets[0];
            if pv_get_bit(&pauli.z, q) {
                Complex64::new(0.0, 1.0)
            } else {
                Complex64::new(1.0, 0.0)
            }
        }
        Gate::SXdg => {
            let q = targets[0];
            if pv_get_bit(&pauli.z, q) {
                Complex64::new(0.0, -1.0)
            } else {
                Complex64::new(1.0, 0.0)
            }
        }
        Gate::Cz => {
            let q0 = targets[0];
            let q1 = targets[1];
            if pv_get_bit(&pauli.x, q0) && pv_get_bit(&pauli.x, q1) {
                Complex64::new(-1.0, 0.0)
            } else {
                Complex64::new(1.0, 0.0)
            }
        }
        _ => Complex64::new(1.0, 0.0),
    }
}

struct WeightedPauliSum {
    terms: HashMap<PauliVec, Complex64>,
    scratch: Vec<(PauliVec, Complex64)>,
}

impl WeightedPauliSum {
    fn new() -> Self {
        Self {
            terms: HashMap::new(),
            scratch: Vec::new(),
        }
    }

    fn insert(&mut self, pauli: PauliVec, coeff: Complex64) {
        let entry = self.terms.entry(pauli).or_insert(Complex64::new(0.0, 0.0));
        *entry += coeff;
    }

    fn conjugate_all_backward_phased(&mut self, gate: &Gate, targets: &[usize]) {
        // Clifford conjugation is a bijection on Pauli strings: the symplectic
        // action on the (x, z) bits is invertible, so distinct keys map to
        // distinct keys and no coefficient merging occurs. Reuse `scratch` to
        // avoid a per-gate heap allocation, drain into it (which keeps the
        // map's bucket capacity), then re-insert directly without the
        // accumulate path.
        self.scratch.clear();
        self.scratch.extend(self.terms.drain());
        for (mut pauli, coeff) in self.scratch.drain(..) {
            let phase = clifford_conjugation_phase(gate, targets, &pauli);
            propagate_backward(&mut pauli, gate, targets);
            self.terms.insert(pauli, coeff * phase);
        }
    }

    fn branch_t_deterministic(&mut self, qubit: usize, is_dagger: bool) {
        let old_terms: Vec<(PauliVec, Complex64)> = self.terms.drain().collect();
        let inv_sqrt2 = 1.0 / SQRT_2;
        let keep_coeff = Complex64::new(inv_sqrt2, 0.0);
        let flip_coeff = if is_dagger {
            Complex64::new(0.0, inv_sqrt2)
        } else {
            Complex64::new(0.0, -inv_sqrt2)
        };

        for (pauli, coeff) in old_terms {
            if !pauli.has_x_or_y(qubit) {
                self.insert(pauli, coeff);
                continue;
            }

            let pauli_keep = pauli.clone();
            let mut pauli_flip = pauli;
            flip_bit(&mut pauli_flip.z, qubit);

            self.insert(pauli_keep, coeff * keep_coeff);
            self.insert(pauli_flip, coeff * flip_coeff);
        }
    }

    fn truncate(&mut self, epsilon: f64) -> f64 {
        let mut discarded = 0.0;
        self.terms.retain(|_, coeff| {
            if coeff.norm() < epsilon {
                discarded += coeff.norm();
                false
            } else {
                true
            }
        });
        discarded
    }

    fn diagonal_expectation(&self) -> f64 {
        let mut sum = Complex64::new(0.0, 0.0);
        for (pauli, coeff) in &self.terms {
            if pauli.is_diagonal() {
                sum += coeff;
            }
        }
        sum.re
    }
}

/// Per-qubit `⟨Z_q⟩` values from a deterministic sparse Pauli dynamics run.
pub struct SpdResult {
    /// `⟨Z_q⟩` per qubit; exact when nothing was truncated.
    pub expectations: Vec<f64>,
    /// T plus Tdg gates in the circuit.
    pub t_count: usize,
    /// Peak size of the weighted Pauli sum across all per-qubit runs, not the
    /// truncation budget passed in.
    pub max_terms: usize,
    /// Total coefficient magnitude discarded by truncation (0 for exact runs).
    pub total_discarded: f64,
}

/// Deterministic SPD on a joint Pauli observable.
///
/// Starts with the single weighted term `(observable, 1.0)`, backward-
/// propagates through every gate, branches each T into two terms with
/// `α / β` coefficients, and truncates terms whose magnitude falls below
/// `epsilon` whenever the sum exceeds `max_terms`. Returns
/// `⟨0^n| U† P U |0^n⟩` as the sum of remaining diagonal-term
/// coefficients.
pub fn run_spd_observable(
    circuit: &Circuit,
    observable: &[PauliTerm],
    epsilon: f64,
    max_terms: usize,
) -> Result<SpdObservableResult> {
    validate_clifford_t_unitary(circuit, "SPD observable")?;
    let n = circuit.num_qubits;
    let t_count = count_t_gates(circuit);
    let (obs, obs_coeff) = pauli_vec_from_terms(n, observable)?;

    let mut sum = WeightedPauliSum::new();
    sum.insert(obs, obs_coeff);
    let mut peak_terms = sum.terms.len();
    let mut total_discarded = 0.0;

    for inst in circuit.instructions.iter().rev() {
        if let Instruction::Gate { gate, targets } = inst {
            match gate {
                Gate::T => sum.branch_t_deterministic(targets[0], false),
                Gate::Tdg => sum.branch_t_deterministic(targets[0], true),
                _ => sum.conjugate_all_backward_phased(gate, targets),
            }
        }
        if max_terms > 0 && sum.terms.len() > max_terms {
            total_discarded += sum.truncate(epsilon);
        }
        check_spd_term_ceiling(sum.terms.len())?;
        if sum.terms.len() > peak_terms {
            peak_terms = sum.terms.len();
        }
    }

    if epsilon > 0.0 {
        total_discarded += sum.truncate(epsilon);
    }

    Ok(SpdObservableResult {
        mean: sum.diagonal_expectation(),
        t_count,
        peak_terms,
        total_discarded,
    })
}

/// Compute `⟨Z_q⟩` for every qubit of a unitary Clifford+T circuit via
/// deterministic sparse Pauli dynamics.
///
/// Each qubit's `Z_q` propagates backward as a weighted Pauli sum; T gates
/// double the in-support terms. When the sum exceeds `max_terms`, terms with
/// coefficient magnitude below `epsilon` are dropped. `max_terms == 0`
/// disables truncation: exact, but the sum must stay under a hard term
/// ceiling. See [`run_spd_observable`] for a single joint observable.
pub fn run_spd(circuit: &Circuit, epsilon: f64, max_terms: usize) -> Result<SpdResult> {
    validate_clifford_t_unitary(circuit, "SPD")?;
    let n = circuit.num_qubits;
    let num_words = n.div_ceil(64);
    let t_count = count_t_gates(circuit);

    let mut expectations = Vec::with_capacity(n);
    let mut peak_terms = 0usize;
    let mut total_discarded = 0.0;

    for q in 0..n {
        let mut sum = WeightedPauliSum::new();
        sum.insert(PauliVec::z_on_qubit(num_words, q), Complex64::new(1.0, 0.0));

        for inst in circuit.instructions.iter().rev() {
            if let Instruction::Gate { gate, targets } = inst {
                match gate {
                    Gate::T => sum.branch_t_deterministic(targets[0], false),
                    Gate::Tdg => sum.branch_t_deterministic(targets[0], true),
                    _ => sum.conjugate_all_backward_phased(gate, targets),
                }
            }

            if max_terms > 0 && sum.terms.len() > max_terms {
                total_discarded += sum.truncate(epsilon);
            }
            check_spd_term_ceiling(sum.terms.len())?;

            if sum.terms.len() > peak_terms {
                peak_terms = sum.terms.len();
            }
        }

        if epsilon > 0.0 {
            total_discarded += sum.truncate(epsilon);
        }

        expectations.push(sum.diagonal_expectation());
    }

    Ok(SpdResult {
        expectations,
        t_count,
        max_terms: peak_terms,
        total_discarded,
    })
}

/// Inverse light cone of a Pauli observable under a circuit, computed
/// conservatively by gate-graph reachability.
///
/// Returns, for each instruction index in `circuit.instructions`, whether the
/// gate at that index can affect the backward-propagated observable. A gate is
/// in the cone if its support intersects the current cone-qubit set when the
/// circuit is traversed in reverse from the observable.
///
/// Exactness: a gate whose target set is disjoint from the cone-qubit set at
/// its backward-traversal depth conjugates the propagated observable trivially
/// (`U_k^dag P_k U_k = P_k` because the support of `P_k` is contained in the
/// cone set at depth `k`, and `U_k` acts as identity outside its targets).
/// Removing those gates from the backward pass is therefore exact.
pub fn inverse_light_cone(circuit: &Circuit, observable: &[PauliTerm]) -> Vec<bool> {
    let mut cone: std::collections::HashSet<usize> = observable.iter().map(|t| t.qubit).collect();
    let n_inst = circuit.instructions.len();
    let mut keep = vec![false; n_inst];

    for (idx, inst) in circuit.instructions.iter().enumerate().rev() {
        let targets: &[usize] = match inst {
            Instruction::Gate { targets, .. } => targets,
            _ => continue,
        };
        let touches = targets.iter().any(|q| cone.contains(q));
        if touches {
            keep[idx] = true;
            for &q in targets {
                cone.insert(q);
            }
        }
    }
    keep
}

/// SPD on a joint Pauli observable, restricted to the inverse light cone.
///
/// Identical in result to `run_spd_observable` for any Clifford+T circuit, but
/// skips gates whose support is disjoint from the propagated observable's
/// causal cone. For QEC syndrome and detector observables with bounded
/// spatial support, this turns the SPD cliff from a function of total T-count
/// into a function of in-cone T-count.
pub fn run_spd_observable_light_cone(
    circuit: &Circuit,
    observable: &[PauliTerm],
    epsilon: f64,
    max_terms: usize,
) -> Result<SpdObservableResult> {
    validate_clifford_t_unitary(circuit, "light-cone SPD observable")?;
    let n = circuit.num_qubits;
    let t_count = count_t_gates(circuit);
    let (obs, obs_coeff) = pauli_vec_from_terms(n, observable)?;
    let keep = inverse_light_cone(circuit, observable);

    let mut sum = WeightedPauliSum::new();
    sum.insert(obs, obs_coeff);
    let mut peak_terms = sum.terms.len();
    let mut total_discarded = 0.0;

    for (idx, inst) in circuit.instructions.iter().enumerate().rev() {
        if !keep[idx] {
            continue;
        }
        if let Instruction::Gate { gate, targets } = inst {
            match gate {
                Gate::T => sum.branch_t_deterministic(targets[0], false),
                Gate::Tdg => sum.branch_t_deterministic(targets[0], true),
                _ => sum.conjugate_all_backward_phased(gate, targets),
            }
        }
        if max_terms > 0 && sum.terms.len() > max_terms {
            total_discarded += sum.truncate(epsilon);
        }
        check_spd_term_ceiling(sum.terms.len())?;
        if sum.terms.len() > peak_terms {
            peak_terms = sum.terms.len();
        }
    }

    if epsilon > 0.0 {
        total_discarded += sum.truncate(epsilon);
    }

    Ok(SpdObservableResult {
        mean: sum.diagonal_expectation(),
        t_count,
        peak_terms,
        total_discarded,
    })
}

#[cfg(test)]
fn spd_to_probabilities(result: &SpdResult) -> Vec<f64> {
    super::expectations_to_marginals(&result.expectations)
        .into_iter()
        .flat_map(|(p0, p1)| [p0, p1])
        .collect()
}

#[cfg(test)]
#[path = "unified_pauli_tests.rs"]
mod tests;