"""Stabilizer-code quantum error correction utilities: generic Pauli-string
commutation/syndrome primitives, and an erasure-aware decoder.

Promoted from Dense-Evolution-Discovery's Steane [[7,1,3]] code
investigation (scripts/steane_code_block6_erasure_conversion.py), where a
Steane-specific version of `erasure_aware_decode` was first built and
verified: on shots with exactly 2 simultaneous heralded erasures, it
achieved exactly 0 decoding failures across 94-12,469 such shots at every
tested physical error rate (0 failures out of >60,000 double-erasure
shots total, 40,000 trials x 10 p-values), versus ~25% failure for a
standard syndrome-only decoder blind to the erasure locations -- a clean
confirmation of the real erasure-correction bound below. This version is
code-agnostic (works from any stabilizer generator list, not a
hand-built Steane-specific table), so it moved here instead of staying
Discovery-repo-specific research code.

Erasure-aware decoding exploits a real, foundational fact: Grassl, Beth &
Pellizzari, "Codes for the quantum erasure channel", Phys. Rev. A 56, 33
(1997) -- a distance-d stabilizer code can correct up to (d-1) ERASURES
(known-location errors, e.g. a heralded lost photon in a dual-rail
photonic qubit), versus only floor((d-1)/2) arbitrary (unlocated)
errors. Erasure location information is worth roughly twice as much as
an ordinary syndrome bit, because knowing WHERE the error is removes
exactly the ambiguity a blind syndrome-only decoder has to guess at.
"""
import itertools
from typing import Optional, Sequence

_ANTICOMMUTING_PAIRS = {
    frozenset(('X', 'Z')), frozenset(('X', 'Y')), frozenset(('Y', 'Z')),
}
PAULIS = ('I', 'X', 'Y', 'Z')


def pauli_commutes(p1: str, p2: str) -> bool:
    """Whether two equal-length Pauli strings (each character in IXYZ, no
    global phase) commute -- the standard symplectic rule: they commute
    iff the number of qubit positions where the local single-qubit
    Paulis anticommute (X/Z, X/Y, or Y/Z, in either order; I commutes
    with everything) is EVEN.

    >>> pauli_commutes('XX', 'ZZ')   # X,Z anticommute at both qubits -> 2 (even) -> commute
    True
    >>> pauli_commutes('XI', 'ZI')   # X,Z anticommute at 1 qubit -> 1 (odd) -> anticommute
    False
    """
    if len(p1) != len(p2):
        raise ValueError(f"Pauli strings must be equal length: {len(p1)} != {len(p2)}")
    n_anticommuting = sum(
        1 for a, b in zip(p1, p2)
        if a != 'I' and b != 'I' and a != b and frozenset((a, b)) in _ANTICOMMUTING_PAIRS
    )
    return n_anticommuting % 2 == 0


def compute_syndrome(pauli_error: str, stabilizers: Sequence[str]) -> tuple:
    """The syndrome (one bit per stabilizer generator, 1 = anticommutes /
    detected, 0 = commutes / undetected) a given Pauli error string would
    produce against `stabilizers` (a list of equal-length Pauli strings,
    the code's stabilizer generators -- X-type, Z-type, or mixed; this
    function doesn't assume a CSS structure)."""
    return tuple(0 if pauli_commutes(pauli_error, g) else 1 for g in stabilizers)


def _pauli_string(n_qubits: int, assignment: dict) -> str:
    chars = ['I'] * n_qubits
    for q, p in assignment.items():
        chars[q] = p
    return ''.join(chars)


def erasure_aware_decode(
    observed_syndrome: tuple,
    heralded_qubits: Sequence[int],
    n_qubits: int,
    stabilizers: Sequence[str],
) -> Optional[str]:
    """Erasure-aware decoder for any stabilizer code. Given the observed
    syndrome and a list of qubits KNOWN to have been erased (e.g. a
    heralded photon-loss event on a dual-rail-encoded qubit), brute-forces
    every Pauli assignment (I/X/Y/Z, 4**len(heralded_qubits) combinations)
    on just the heralded qubits and returns the unique full-length Pauli
    string reproducing `observed_syndrome` exactly.

    Returns `None` -- not a guess -- when there are zero heralded qubits,
    when the observed syndrome is not explained by any assignment on the
    heralded qubits alone, or when more than one assignment explains it
    (ambiguous). Both `None` cases mean: fall back to a standard
    syndrome-only decoder, or treat as a detected-but-uncorrectable
    event -- this function will not silently return a wrong-but-plausible
    correction. The number of heralded qubits this can actually resolve
    unambiguously is bounded by the code's real distance (Grassl, Beth &
    Pellizzari 1997: up to d-1 erasures) -- that bound emerges naturally
    from the brute-force search itself (more heralded qubits than the
    code can resolve typically yields zero or multiple matches), it is
    not hard-coded here.

    Cost is 4**len(heralded_qubits) syndrome computations, each O(n_qubits
    * len(stabilizers)) -- fine for the small numbers of simultaneous
    erasures a real per-shot noise rate produces (verified up to 2 in the
    original Steane investigation; tractable up to 4-5 for most small
    codes before it's worth switching to a smarter search).
    """
    if not heralded_qubits:
        return None

    matches = []
    for assignment_paulis in itertools.product(PAULIS, repeat=len(heralded_qubits)):
        assignment = dict(zip(heralded_qubits, assignment_paulis))
        candidate = _pauli_string(n_qubits, assignment)
        if compute_syndrome(candidate, stabilizers) == tuple(observed_syndrome):
            matches.append(candidate)

    if len(matches) != 1:
        return None
    return matches[0]