"""Stim-based stabilizer verifier.

Vendored / adapted from uw-math-ai/quantum-ai's tools/check_stabilizers.py.
The core check uses Stim's TableauSimulator + peek_observable_expectation,
which is exact and polynomial in the number of qubits for Clifford circuits.
"""
from __future__ import annotations

import re

import stim


def preprocess_stim_text(raw: str) -> str:
    if raw is None:
        return ""
    text = raw.strip().replace("\\n", "\n")
    lines = []
    for line in text.splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        line = re.sub(r"\s+", " ", line)
        lines.append(line)
    return "\n".join(lines) + ("\n" if lines else "")


def check_stabilizers(
    circuit_text: str, stabilizers: list[str], n_qubits: int | None = None
) -> dict[str, bool]:
    """Return {stabilizer: preserved?} for a circuit applied to |0...0>.

    If n_qubits is provided, ensures the simulator has at least that many qubits
    by prepending an I gate on qubit n-1. This matters because Stim's simulator
    only allocates qubits that are touched by the circuit.
    """
    text = preprocess_stim_text(circuit_text)
    if n_qubits is not None and n_qubits > 0:
        text = f"I {n_qubits - 1}\n" + text
    if not text:
        text = "I 0\n"

    circ = stim.Circuit(text)
    sim = stim.TableauSimulator()
    sim.do(circ)

    width = max(circ.num_qubits, n_qubits or 0)
    results: dict[str, bool] = {}
    for stab in stabilizers:
        padded = stab + "I" * max(0, width - len(stab))
        pauli = stim.PauliString(padded)
        expectation = sim.peek_observable_expectation(pauli)
        results[stab] = expectation > 0
    return results


def match_fraction(
    circuit_text: str, stabilizers: list[str], n_qubits: int
) -> tuple[float, dict[str, bool]]:
    res = check_stabilizers(circuit_text, stabilizers, n_qubits=n_qubits)
    if not res:
        return 0.0, res
    frac = sum(1 for v in res.values() if v) / len(res)
    return frac, res