stabilizer_qec 1.5.0

use super::{probability, BitTable, DetectorErrorModel, Error, Result};
use crate::batch::{StreamedWindows, belief_shots, bposd_shots, match_shots, streamed_window_shots, union_find_shots, window_info, window_shots};
use crate::dem_decoder::{DecodeError, DemDecoder};
use crate::sparse::{Correlations, Scratch, SparseGraph};
use crate::window::{Mode, Model, WindowDecoder};

/// A decoder's answer for one shot.
#[derive(Clone, Copy, Debug, PartialEq)]
#[non_exhaustive]
pub struct Prediction {
    /// The observables predicted flipped, observable `k` in bit `k`.
    pub observables: u64,
    /// The correction's weight (Σ ln((1 − p)/p) over its edges), where the decoder matched.
    pub weight: Option<f64>,
    /// Whether belief propagation converged on its own, where the decoder ran it.
    pub bp_converged: Option<bool>,
}

impl Prediction {
    /// Whether observable `k` is predicted flipped.
    pub fn flips(&self, k: usize) -> bool {
        k < 64 && (self.observables >> k) & 1 == 1
    }
}

fn decode_error(e: DecodeError) -> Error {
    Error::new(match e {
        DecodeError::Unmatchable => "no matching explains the detection events (a defect with no path to a partner or the boundary)".to_string(),
        other => format!("the matcher declined the shot ({other:?})"),
    })
}

/// Check a batch's width and say which shot failed, if any.
fn check_width(shots: &BitTable, num_detectors: usize) -> Result<()> {
    if shots.num_bits() != num_detectors {
        return Err(Error::new(format!("{} bits per shot, not the model's {num_detectors} detectors", shots.num_bits())));
    }
    Ok(())
}

fn check_defects(defects: &[u32], num_detectors: usize) -> Result<()> {
    match defects.iter().find(|&&d| d as usize >= num_detectors) {
        Some(d) => Err(Error::new(format!("detector {d} is beyond the model's {num_detectors}"))),
        None => Ok(()),
    }
}

/// Exact minimum-weight perfect matching on a decomposed detector error model, by sparse
/// blossom, as PyMatching decodes: edges weighted ln((1 − p)/p), the observables of the lightest
/// correction returned. Correlated matching adds PyMatching 2.4's second pass.
///
/// `Send + Sync`: share one across threads, or let [`Matching::decode_batch`] cut a batch across
/// them.
///
/// ```
/// use stabilizer_qec::{DetectorErrorModel, Matching};
///
/// let dem: DetectorErrorModel = "error(0.1) D0\nerror(0.1) D0 D1 L0\nerror(0.1) D1".parse()?;
/// let m = Matching::new(&dem)?;
/// assert!(m.decode(&[0, 1])?.flips(0));
/// # Ok::<(), stabilizer_qec::Error>(())
/// ```
pub struct Matching {
    graph: SparseGraph,
    corr: Option<Correlations>,
    num_detectors: usize,
    num_observables: usize,
}

impl Matching {
    /// Plain matching. The model's faults must flip at most two detectors each, or be
    /// decomposed into such pieces.
    pub fn new(dem: &DetectorErrorModel) -> Result<Matching> {
        Matching::build(dem, false)
    }

    /// Correlated matching (PyMatching's `enable_correlations=True`).
    pub fn with_correlations(dem: &DetectorErrorModel) -> Result<Matching> {
        Matching::build(dem, true)
    }

    fn build(dem: &DetectorErrorModel, correlated: bool) -> Result<Matching> {
        let (graph, corr) = DemDecoder::new(dem.flat()?)?.into_parts(correlated);
        Ok(Matching { graph, corr, num_detectors: dem.num_detectors(), num_observables: dem.num_observables() })
    }

    /// Detectors in the model.
    pub fn num_detectors(&self) -> usize {
        self.num_detectors
    }

    /// Observables in the model.
    pub fn num_observables(&self) -> usize {
        self.num_observables
    }

    /// One shot, given the detectors that fired.
    pub fn decode(&self, defects: &[u32]) -> Result<Prediction> {
        check_defects(defects, self.num_detectors)?;
        let mut sorted = defects.to_vec();
        sorted.sort_unstable();
        sorted.dedup();
        let mut scratch = Scratch::new(&self.graph);
        let p = match &self.corr {
            Some(corr) => self.graph.decode_correlated(corr, &mut scratch, &sorted),
            None => self.graph.decode(&mut scratch, &sorted),
        }
        .map_err(decode_error)?;
        Ok(Prediction { observables: p.observables, weight: Some(p.weight), bp_converged: None })
    }

    /// A batch of shots (one row per shot, one bit per detector), across `threads` threads
    /// (`0` is every core). A shot no matching explains is an error naming it.
    pub fn decode_batch(&self, shots: &BitTable, threads: usize) -> Result<Vec<Prediction>> {
        check_width(shots, self.num_detectors)?;
        let out = match_shots(&self.graph, self.corr.as_ref(), shots.as_bytes(), self.num_detectors, shots.num_rows(), threads);
        if let Some(s) = out.iter().position(|(_, w)| w.is_nan()) {
            return Err(Error::new(format!("shot {s}: no matching explains its detection events")));
        }
        Ok(out.into_iter().map(|(o, w)| Prediction { observables: o, weight: Some(w), bp_converged: None }).collect())
    }
}

/// How belief propagation updates its messages.
#[derive(Clone, Copy, Debug, PartialEq)]
#[non_exhaustive]
pub enum BpMethod {
    /// Product-sum (sum-product) BP.
    ProductSum,
    /// Min-sum BP, its check messages scaled by `scaling_factor` (`0.0`: `ldpc`'s adaptive
    /// scaling, 1 − 2^−iteration).
    MinimumSum {
        /// The scaling factor.
        scaling_factor: f64,
    },
}

impl BpMethod {
    fn engine(self) -> crate::bp::Method {
        match self {
            BpMethod::ProductSum => crate::bp::Method::ProductSum,
            BpMethod::MinimumSum { scaling_factor } => crate::bp::Method::MinSum { scale: scaling_factor },
        }
    }
}

/// Belief-propagation settings.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct BpOptions {
    max_iter: usize,
    method: BpMethod,
}

impl BpOptions {
    /// `max_iter` iterations of `method`.
    pub fn new(max_iter: usize, method: BpMethod) -> BpOptions {
        BpOptions { max_iter, method }
    }
}

/// How BP+OSD's ordered-statistics step searches when BP does not converge.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub enum OsdMethod {
    /// OSD-0: the most likely information set alone.
    Osd0,
    /// Exhaustive search over the given order's least reliable columns.
    Exhaustive(usize),
    /// Combination sweep of the given order (`ldpc`'s `osd_cs`).
    CombinationSweep(usize),
}

impl OsdMethod {
    fn engine(self) -> crate::osd::OsdMethod {
        match self {
            OsdMethod::Osd0 => crate::osd::OsdMethod::Osd0,
            OsdMethod::Exhaustive(k) => crate::osd::OsdMethod::Exhaustive(k),
            OsdMethod::CombinationSweep(k) => crate::osd::OsdMethod::CombinationSweep(k),
        }
    }
}

/// Weighted union-find decoding (Delfosse and Nickerson, with Huang, Newman and Brown's
/// weighted growth) on the same graph matching uses: clusters around the detection events grow
/// edge weight by edge weight until their events can pair inside them, and a spanning tree of
/// each is peeled for the correction. Linear time in practice, its logical error rate a little
/// above matching's; the standard baseline (this engine's matcher, near-linear too, is about
/// twice as fast). The model's faults must flip at most two detectors each, or be decomposed.
///
/// ```
/// use stabilizer_qec::{memory_circuit, Basis, DemOptions, Noise, SurfaceCode, UnionFind};
///
/// let c = memory_circuit(SurfaceCode::Rotated, 5, 5, Noise::Sd6 { p: 0.003 }, Basis::Z)?;
/// let dem = c.detector_error_model(&DemOptions::new().decompose_errors(true))?;
/// let samples = c.detector_sampler(1)?.sample(2000, 0);
/// let predictions = UnionFind::new(&dem)?.decode_batch(&samples.detectors, 0)?;
/// let wrong = (0..2000).filter(|&s| predictions[s].flips(0) != samples.observables.get(s, 0)).count();
/// assert!(wrong < 100);
/// # Ok::<(), stabilizer_qec::Error>(())
/// ```
pub struct UnionFind {
    graph: SparseGraph,
    num_detectors: usize,
    num_observables: usize,
}

impl UnionFind {
    /// The decoder for a decomposed model.
    pub fn new(dem: &DetectorErrorModel) -> Result<UnionFind> {
        let (graph, _) = DemDecoder::new(dem.flat()?)?.into_parts(false);
        Ok(UnionFind { graph, num_detectors: dem.num_detectors(), num_observables: dem.num_observables() })
    }

    /// Detectors in the model.
    pub fn num_detectors(&self) -> usize {
        self.num_detectors
    }

    /// Observables in the model.
    pub fn num_observables(&self) -> usize {
        self.num_observables
    }

    /// One shot, given the detectors that fired.
    pub fn decode(&self, defects: &[u32]) -> Result<Prediction> {
        check_defects(defects, self.num_detectors)?;
        let mut sorted = defects.to_vec();
        sorted.sort_unstable();
        sorted.dedup();
        let mut scratch = self.graph.union_find_scratch();
        let o = self
            .graph
            .decode_union_find(&mut scratch, &sorted)
            .ok_or_else(|| Error::new("no correction explains the detection events (an odd component with no boundary)"))?;
        Ok(Prediction { observables: o, weight: None, bp_converged: None })
    }

    /// A batch of shots (one row per shot, one bit per detector), across `threads` threads
    /// (`0` is every core). A shot no correction explains is an error naming it.
    pub fn decode_batch(&self, shots: &BitTable, threads: usize) -> Result<Vec<Prediction>> {
        check_width(shots, self.num_detectors)?;
        let out = union_find_shots(&self.graph, shots.as_bytes(), self.num_detectors, shots.num_rows(), threads);
        out.into_iter()
            .enumerate()
            .map(|(s, o)| match o {
                Some(observables) => Ok(Prediction { observables, weight: None, bp_converged: None }),
                None => Err(Error::new(format!("shot {s}: no correction explains its detection events"))),
            })
            .collect()
    }
}

/// Belief-matching (Higgott et al. 2023), as the `beliefmatching` package decodes: BP on the
/// model's hypergraph; where BP's own correction explains the shot it is used, otherwise
/// matching on weights −ln p from BP's posteriors. The model must be decomposed.
pub struct BeliefMatching {
    inner: crate::belief::BeliefMatching,
    num_detectors: usize,
}

impl BeliefMatching {
    /// Belief-matching with the given BP settings (the package's defaults are 20 iterations of
    /// product-sum).
    pub fn new(dem: &DetectorErrorModel, bp: BpOptions) -> Result<BeliefMatching> {
        let inner = crate::belief::BeliefMatching::from_dem(dem.flat()?, bp.method.engine(), bp.max_iter)?;
        Ok(BeliefMatching { inner, num_detectors: dem.num_detectors() })
    }

    /// One shot, given the detectors that fired.
    pub fn decode(&self, defects: &[u32]) -> Result<Prediction> {
        check_defects(defects, self.num_detectors)?;
        let mut sorted = defects.to_vec();
        sorted.sort_unstable();
        sorted.dedup();
        let o = self.inner.decode(&sorted, &mut self.inner.work()).map_err(decode_error)?;
        Ok(Prediction { observables: o.observables, weight: (!o.converged).then_some(o.weight), bp_converged: Some(o.converged) })
    }

    /// A batch of shots across `threads` threads (`0` is every core).
    pub fn decode_batch(&self, shots: &BitTable, threads: usize) -> Result<Vec<Prediction>> {
        check_width(shots, self.num_detectors)?;
        let out = belief_shots(&self.inner, shots.as_bytes(), self.num_detectors, shots.num_rows(), threads);
        if let Some(s) = out.iter().position(|x| x.2 == 2) {
            return Err(Error::new(format!("shot {s}: no correction explains its detection events")));
        }
        Ok(out
            .into_iter()
            .map(|(o, w, c)| Prediction { observables: o, weight: (c == 0).then_some(w), bp_converged: Some(c == 1) })
            .collect())
    }
}

/// BP+OSD on a detector error model (Roffe et al.): each fault a column with its prior,
/// corrections equal to `ldpc`'s `BpOsdDecoder` but for ties among equally likely columns. For
/// codes whose faults flip three or more detectors, such as the bivariate bicycle codes: give it
/// the undecomposed model.
pub struct BpOsd {
    inner: crate::osd::BpOsd,
    observables: Vec<u64>,
    num_detectors: usize,
}

impl BpOsd {
    /// BP+OSD with the given BP settings and OSD search.
    pub fn new(dem: &DetectorErrorModel, bp: BpOptions, osd: OsdMethod) -> Result<BpOsd> {
        let d = dem.flat()?;
        let columns: Vec<Vec<u32>> = d.mechanisms.iter().map(|m| m.detectors.clone()).collect();
        let priors: Vec<f64> = d.mechanisms.iter().map(|m| m.p).collect();
        let inner = crate::osd::BpOsd::new(d.num_detectors, columns, &priors, bp.method.engine(), bp.max_iter, osd.engine())?;
        Ok(BpOsd { inner, observables: d.mechanisms.iter().map(|m| m.observables).collect(), num_detectors: d.num_detectors })
    }

    /// One shot, given the detectors that fired.
    pub fn decode(&self, defects: &[u32]) -> Result<Prediction> {
        check_defects(defects, self.num_detectors)?;
        let mut row = BitTable::zeros(1, self.num_detectors);
        for &d in defects {
            row.set(0, d as usize, true);
        }
        Ok(self.decode_batch(&row, 1)?.remove(0))
    }

    /// A batch of shots across `threads` threads (`0` is every core).
    pub fn decode_batch(&self, shots: &BitTable, threads: usize) -> Result<Vec<Prediction>> {
        check_width(shots, self.num_detectors)?;
        let out = bposd_shots(&self.inner, &self.observables, shots.as_bytes(), self.num_detectors, shots.num_rows(), threads);
        Ok(out.into_iter().map(|(o, c)| Prediction { observables: o, weight: None, bp_converged: Some(c == 1) }).collect())
    }
}

/// The order a window decoder runs its windows in.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub enum WindowMode {
    /// One window after another, in time.
    Sliding,
    /// Alternate windows at once, then the gaps between them.
    Parallel,
}

/// A window decoder's schedule.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct WindowOptions {
    commit: usize,
    buffer: usize,
    mode: WindowMode,
    correlations: bool,
    template: Option<bool>,
}

impl Window {
    pub(crate) fn from_info((a, b, c0, c1, phase, waits_for): crate::batch::WindowInfo) -> Window {
        Window { first_layer: a, end_layer: b, commit_start: c0, commit_end: c1, phase, waits_for }
    }
}

impl WindowOptions {
    pub(crate) fn parts(self) -> (usize, usize, WindowMode, bool) {
        (self.commit, self.buffer, self.mode, self.correlations)
    }

    /// Windows that commit `commit` rounds with `buffer` rounds on either side, run in `mode`,
    /// plain matching.
    pub fn new(commit: usize, buffer: usize, mode: WindowMode) -> WindowOptions {
        WindowOptions { commit, buffer, mode, correlations: false, template: None }
    }

    /// Where the windows' graphs come from. `Some(true)`: a short template of the model's
    /// longest loop, whose middle windows serve every window of the long model away from its
    /// ends, shifted (the decoder never unrolls the model: a memory of millions of rounds
    /// decodes); `Some(false)`: the whole model, unrolled; `None` (the default): the template
    /// for a model of more than a million faults that has one, else the whole model. Both give
    /// the same predictions.
    pub fn with_template(mut self, template: Option<bool>) -> WindowOptions {
        self.template = template;
        self
    }

    /// Correlated matching in each window.
    pub fn with_correlations(mut self, yes: bool) -> WindowOptions {
        self.correlations = yes;
        self
    }
}

/// A window of a window decoder: it decodes rounds `first_layer..end_layer` and commits the
/// corrections in `commit_start..commit_end`.
#[derive(Clone, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub struct Window {
    /// The first round it decodes.
    pub first_layer: u32,
    /// One past the last round it decodes.
    pub end_layer: u32,
    /// The first round it commits.
    pub commit_start: u32,
    /// One past the last round it commits.
    pub commit_end: u32,
    /// The phase it runs in (phase 0 first).
    pub phase: usize,
    /// The windows whose commits it waits for.
    pub waits_for: Vec<usize>,
}

/// Window decoding (Skoric et al.; Tan et al.): the model cut into overlapping windows of rounds,
/// each matched alone and committing only its middle. The model's detectors need a time
/// coordinate (their last), as Stim's generated circuits give them.
pub struct WindowMatching {
    inner: Windows,
    correlations: bool,
    num_detectors: usize,
}

enum Windows {
    Whole(Box<WindowDecoder>),
    Streamed(Box<StreamedWindows>),
}

impl WindowMatching {
    /// Window decoding with the given schedule.
    pub fn new(dem: &DetectorErrorModel, options: WindowOptions) -> Result<WindowMatching> {
        let mode = match options.mode {
            WindowMode::Sliding => Mode::Sliding,
            WindowMode::Parallel => Mode::Parallel,
        };
        let streamed = || StreamedWindows::new(dem.program(), options.commit, options.buffer, mode).map(Box::new);
        let inner = match options.template {
            Some(true) => Windows::Streamed(streamed()?),
            None if dem.num_errors() > 1_000_000 => match streamed() {
                Ok(s) => Windows::Streamed(s),
                Err(_) => Windows::Whole(Box::new(WindowDecoder::new(Model::new(dem.flat()?)?, options.commit, options.buffer, mode)?)),
            },
            _ => Windows::Whole(Box::new(WindowDecoder::new(Model::new(dem.flat()?)?, options.commit, options.buffer, mode)?)),
        };
        Ok(WindowMatching { inner, correlations: options.correlations, num_detectors: dem.num_detectors() })
    }

    /// The windows, in the order they are planned.
    pub fn windows(&self) -> Vec<Window> {
        let info = match &self.inner {
            Windows::Whole(wd) => window_info(wd),
            Windows::Streamed(sw) => sw.windows(),
        };
        info.into_iter().map(Window::from_info).collect()
    }

    /// A batch of shots across `threads` threads (`0` is every core).
    pub fn decode_batch(&self, shots: &BitTable, threads: usize) -> Result<Vec<Prediction>> {
        check_width(shots, self.num_detectors)?;
        let out = match &self.inner {
            Windows::Whole(wd) => window_shots(wd, shots.as_bytes(), self.num_detectors, shots.num_rows(), self.correlations, threads, false),
            Windows::Streamed(sw) => streamed_window_shots(sw, shots.as_bytes(), self.num_detectors, shots.num_rows(), self.correlations, threads, false),
        };
        if let Some(s) = out.iter().position(|x| x.3) {
            return Err(Error::new(format!("shot {s}: a window found no matching")));
        }
        if out.iter().any(|x| x.1 > 0) {
            return Err(Error::new("internal error: window commits left defects unexplained; please report it"));
        }
        Ok(out.into_iter().map(|x| Prediction { observables: x.0, weight: None, bp_converged: None }).collect())
    }
}

fn check_matrix(num_checks: usize, columns: &[Vec<u32>], priors: &[f64]) -> Result<()> {
    if columns.len() != priors.len() {
        return Err(Error::new(format!("{} columns but {} priors", columns.len(), priors.len())));
    }
    for (j, &p) in priors.iter().enumerate() {
        probability(p, &format!("prior {j}"))?;
    }
    if let Some((j, r)) = columns.iter().enumerate().find_map(|(j, c)| c.iter().find(|&&r| r as usize >= num_checks).map(|r| (j, r))) {
        return Err(Error::new(format!("column {j} touches check {r}, beyond {num_checks}")));
    }
    Ok(())
}

fn syndrome_bytes(syndrome: &[bool], num_checks: usize) -> Result<Vec<u8>> {
    if syndrome.len() != num_checks {
        return Err(Error::new(format!("{} syndrome bits for {num_checks} checks", syndrome.len())));
    }
    Ok(syndrome.iter().map(|&b| u8::from(b)).collect())
}

/// A BP run's result.
#[derive(Clone, Debug, PartialEq)]
#[non_exhaustive]
pub struct BpOutcome {
    /// The hard decision, one per column.
    pub correction: Vec<bool>,
    /// Each column's posterior log-likelihood ratio, ln(P(0)/P(1)).
    pub log_prob_ratios: Vec<f64>,
    /// Whether the hard decision explains the syndrome.
    pub converged: bool,
    /// Iterations run.
    pub iterations: usize,
}

/// Belief propagation on a check matrix given by its columns (the checks each column touches),
/// flooding schedule, its arithmetic reproducing `ldpc`'s `BpDecoder` bit for bit.
pub struct BpDecoder {
    inner: crate::bp::Bp,
    options: BpOptions,
    num_checks: usize,
}

impl BpDecoder {
    /// A decoder for `num_checks` checks and the given columns, each with its prior.
    pub fn new(num_checks: usize, columns: &[Vec<u32>], priors: &[f64], options: BpOptions) -> Result<BpDecoder> {
        check_matrix(num_checks, columns, priors)?;
        Ok(BpDecoder { inner: crate::bp::Bp::new(num_checks, columns, priors)?, options, num_checks })
    }

    /// Decode one syndrome (one bit per check).
    pub fn decode(&self, syndrome: &[bool]) -> Result<BpOutcome> {
        let s = syndrome_bytes(syndrome, self.num_checks)?;
        let mut w = self.inner.work();
        let out = self.inner.decode(&s, self.options.method.engine(), self.options.max_iter, &mut w);
        Ok(BpOutcome {
            correction: w.hard.iter().map(|&b| b != 0).collect(),
            log_prob_ratios: w.llr.clone(),
            converged: out.converged,
            iterations: out.iterations,
        })
    }
}

/// A BP+OSD run's result.
#[derive(Clone, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub struct BpOsdOutcome {
    /// The correction, one per column.
    pub correction: Vec<bool>,
    /// Whether BP converged before OSD.
    pub converged: bool,
    /// BP's iterations.
    pub iterations: usize,
}

/// BP+OSD on a check matrix given by its columns, as `ldpc`'s `BpOsdDecoder`.
pub struct BpOsdDecoder {
    inner: crate::osd::BpOsd,
    num_checks: usize,
}

impl BpOsdDecoder {
    /// A decoder for `num_checks` checks and the given columns, each with its prior.
    pub fn new(num_checks: usize, columns: &[Vec<u32>], priors: &[f64], bp: BpOptions, osd: OsdMethod) -> Result<BpOsdDecoder> {
        check_matrix(num_checks, columns, priors)?;
        let inner = crate::osd::BpOsd::new(num_checks, columns.to_vec(), priors, bp.method.engine(), bp.max_iter, osd.engine())?;
        Ok(BpOsdDecoder { inner, num_checks })
    }

    /// Decode one syndrome (one bit per check).
    pub fn decode(&self, syndrome: &[bool]) -> Result<BpOsdOutcome> {
        let s = syndrome_bytes(syndrome, self.num_checks)?;
        let mut w = self.inner.work();
        let out = self.inner.decode(&s, &mut w);
        Ok(BpOsdOutcome { correction: w.correction.iter().map(|&b| b != 0).collect(), converged: out.converged, iterations: out.iterations })
    }
}