instructions

def add_dummy(
    b: GraphRepresentation, qubit: int, row: float | int | None = None
) -> int:
    """Add a dummy boundary vertex for a qubit."""
    if row is None:
        row = last_row(b, qubit) + 1
    v1 = b.graph.add_vertex(VertexType.BOUNDARY, qubit=qubit, row=row)
    b.last_vertex[qubit] = v1
    return v1

def add_lane(b: GraphRepresentation, qubit: int) -> int:
    """Initialize a qubit lane if it doesn't exist."""
    v1 = b.graph.add_vertex(VertexType.BOUNDARY, qubit=qubit, row=0)
    v2 = b.graph.add_vertex(VertexType.BOUNDARY, qubit=qubit, row=1)
    b.graph.add_edge((v1, v2), b.edge_type.SIMPLE)
    b.first_vertex[qubit] = v1
    b.last_vertex[qubit] = v2
    return v1

def c_nxyz(b: GraphRepresentation, qubit: int) -> None:
    """Period 3 axis cycling gate, sending -X -> Y -> Z -> -X."""
    s_dag(b, qubit)
    sqrt_y_dag(b, qubit)
    b.graph.scalar.add_phase(Fraction(1, 4))

def c_nzyx(b: GraphRepresentation, qubit: int) -> None:
    """Period 3 axis cycling gate, sending -Z -> Y -> X -> -Z."""
    s_dag(b, qubit)
    sqrt_x(b, qubit)
    b.graph.scalar.add_phase(Fraction(-1, 4))

def c_xnyz(b: GraphRepresentation, qubit: int) -> None:
    """Period 3 axis cycling gate, sending X -> -Y -> Z -> X."""
    s(b, qubit)
    h(b, qubit)

def c_xynz(b: GraphRepresentation, qubit: int) -> None:
    """Period 3 axis cycling gate, sending X -> Y -> -Z -> X."""
    s(b, qubit)
    sqrt_y_dag(b, qubit)
    b.graph.scalar.add_phase(Fraction(1, 4))

def c_xyz(b: GraphRepresentation, qubit: int) -> None:
    """Right handed period 3 axis cycling gate, sending X -> Y -> Z -> X."""
    s_dag(b, qubit)
    h(b, qubit)
    b.graph.scalar.add_phase(Fraction(-1, 4))

def c_znyx(b: GraphRepresentation, qubit: int) -> None:
    """Period 3 axis cycling gate, sending Z -> -Y -> X -> Z."""
    s(b, qubit)
    sqrt_x(b, qubit)
    b.graph.scalar.add_phase(Fraction(-1, 4))

def c_zynx(b: GraphRepresentation, qubit: int) -> None:
    """Period 3 axis cycling gate, sending Z -> Y -> -X -> Z."""
    s(b, qubit)
    sqrt_x_dag(b, qubit)
    b.graph.scalar.add_phase(Fraction(1, 4))

def c_zyx(b: GraphRepresentation, qubit: int) -> None:
    """Left handed period 3 axis cycling gate, sending Z -> Y -> X -> Z."""
    h(b, qubit)
    s(b, qubit)
    b.graph.scalar.add_phase(Fraction(1, 4))

def cnot(
    b: GraphRepresentation,
    control: int,
    target: int,
    classically_controlled: list[bool] | None = None,
) -> None:
    """Apply CNOT (controlled-X) gate."""
    _cx_cz(b, True, control, target, classically_controlled)

def correlated_error(
    b: GraphRepresentation,
    qubits: list[int],
    types: list[Literal["X", "Y", "Z"]],
    p: float,
) -> None:
    """Add a correlated error term affecting multiple qubits with given Pauli types."""
    for qubit, type_ in zip(qubits, types, strict=True):
        if type_ == "X" or type_ == "Y":
            _error(b, qubit, VertexType.X, f"c{b.num_correlated_error_bits}")
        if type_ == "Z" or type_ == "Y":
            _error(b, qubit, VertexType.Z, f"c{b.num_correlated_error_bits}")

    b.correlated_error_probs.append(p)
    b.num_correlated_error_bits += 1

def cxswap(b: GraphRepresentation, qubit1: int, qubit2: int) -> None:
    """Apply CX then SWAP."""
    cnot(b, qubit1, qubit2)
    swap(b, qubit1, qubit2)

def cy(
    b: GraphRepresentation,
    control: int,
    target: int,
    classically_controlled: list[bool] | None = None,
) -> None:
    """Apply controlled-Y gate."""
    s_dag(b, target)
    cnot(b, control, target, classically_controlled)
    s(b, target)

def cz(
    b: GraphRepresentation,
    control: int,
    target: int,
    classically_controlled: list[bool] | None = None,
) -> None:
    """Apply controlled-Z gate."""
    _cx_cz(b, False, control, target, classically_controlled)

def czswap(b: GraphRepresentation, qubit1: int, qubit2: int) -> None:
    """Apply CZ then SWAP."""
    cz(b, qubit1, qubit2)
    swap(b, qubit1, qubit2)

def depolarize1(b: GraphRepresentation, qubit: int, p: float) -> None:
    """Apply single-qubit depolarizing channel with total error probability p."""
    pauli_channel_1(b, qubit, p / 3, p / 3, p / 3)

def depolarize2(b: GraphRepresentation, qubit_i: int, qubit_j: int, p: float) -> None:
    """Apply two-qubit depolarizing channel with total error probability p."""
    pauli_channel_2(
        b,
        qubit_i,
        qubit_j,
        p / 15,
        p / 15,
        p / 15,
        p / 15,
        p / 15,
        p / 15,
        p / 15,
        p / 15,
        p / 15,
        p / 15,
        p / 15,
        p / 15,
        p / 15,
        p / 15,
        p / 15,
    )

def detector(b: GraphRepresentation, rec: list[int], *args) -> None:
    """Add detector annotation that XORs the given measurement record bits."""
    row = min({b.graph.row(b.rec[r]) for r in rec}) - 0.5
    d_rows = {b.graph.row(d) for d in b.detectors + b.observables}
    while row in d_rows:
        row += 1
    v0 = b.graph.add_vertex(
        VertexType.X, qubit=-1, row=row, phase=f"det[{len(b.detectors)}]"
    )
    for rec_ in rec:
        b.graph.add_edge((v0, b.rec[rec_]))
    b.detectors.append(v0)

def ensure_lane(b: GraphRepresentation, qubit: int) -> None:
    """Ensure qubit lane exists."""
    if qubit not in b.last_vertex:
        add_lane(b, qubit)

def finalize_correlated_error(b: GraphRepresentation) -> None:
    """Finalize the current correlated error channel.

    1. Rename all "c{i}" phases to "e{num_error_bits + i}" in the graph
    2. Compute and append the 2^k probability array to channel_probs
    3. Increment num_error_bits by k
    4. Reset num_correlated_error_bits to 0 and correlated_error_probs to []
    """
    k = b.num_correlated_error_bits

    if k == 0:
        return

    # Rename "c{i}" phases to "e{num_error_bits + i}" in the graph
    for v in b.graph.vertices():
        phase_vars = b.graph._phaseVars.get(v, set())
        new_phase_vars = set()
        for var in phase_vars:
            if isinstance(var, str) and var.startswith("c"):
                # Extract the bit index from "c{i}"
                bit_idx = int(var[1:])
                new_phase_vars.add(f"e{b.num_error_bits + bit_idx}")
            else:
                new_phase_vars.add(var)
        b.graph._phaseVars[v] = new_phase_vars

    # Compute probability array from conditional probabilities
    probs = correlated_error_probs(b.correlated_error_probs)
    b.channel_probs.append(probs)

    b.num_error_bits += k

    # Reset correlated error state
    b.num_correlated_error_bits = 0
    b.correlated_error_probs = []

def h(b: GraphRepresentation, qubit: int) -> None:
    """Apply Hadamard gate."""
    ensure_lane(b, qubit)
    e = last_edge(b, qubit)
    b.graph.set_edge_type(
        e,
        (
            b.edge_type.HADAMARD
            if b.graph.edge_type(e) == b.edge_type.SIMPLE
            else b.edge_type.SIMPLE
        ),
    )

def h_nxy(b: GraphRepresentation, qubit: int) -> None:
    """Apply Hadamard-like gate that sends -X <-> Y, Z -> -Z."""
    x(b, qubit)
    s_dag(b, qubit)

def h_nxz(b: GraphRepresentation, qubit: int) -> None:
    """Apply Hadamard-like gate that sends -X <-> Z."""
    z(b, qubit)
    sqrt_y_dag(b, qubit)
    b.graph.scalar.add_phase(Fraction(1, 4))

def h_nyz(b: GraphRepresentation, qubit: int) -> None:
    """Apply Hadamard-like gate that sends -Y <-> Z, X -> -X."""
    z(b, qubit)
    sqrt_x(b, qubit)
    b.graph.scalar.add_phase(Fraction(-1, 4))

def h_xy(b: GraphRepresentation, qubit: int) -> None:
    """Apply variant of Hadamard gate that swaps the X and Y axes (instead of X and Z)."""
    x(b, qubit)
    s(b, qubit)
    b.graph.scalar.add_phase(Fraction(-1, 4))

def h_yz(b: GraphRepresentation, qubit: int) -> None:
    """Apply variant of Hadamard gate that swaps the Y and Z axes (instead of X and Z)."""
    sqrt_x(b, qubit)
    z(b, qubit)
    b.graph.scalar.add_phase(Fraction(-1, 4))

def heralded_erase(b: GraphRepresentation, qubit: int, p: float) -> None:
    """Apply heralded erasure channel.

    Special case of heralded_pauli_channel_1 with equal probabilities p/4
    for each of I, X, Y, Z when the channel fires.
    """
    heralded_pauli_channel_1(b, qubit, p / 4, p / 4, p / 4, p / 4)

def heralded_pauli_channel_1(
    b: GraphRepresentation,
    qubit: int,
    pi: float,
    px: float,
    py: float,
    pz: float,
) -> None:
    """Apply heralded single-qubit Pauli channel.

    Records a herald bit into the measurement record. When the channel fires
    (with total probability pi+px+py+pz), the herald is 1 and one of I/X/Y/Z
    is applied. When it doesn't fire, the herald is 0 and nothing happens.
    """
    b.channel_probs.append(heralded_pauli_channel_1_probs(pi, px, py, pz))
    aux = -2
    r(b, aux)
    _error(b, aux, b.vertex_type.X, f"e{b.num_error_bits}")  # herald bit flip
    m(b, aux)
    _error(b, qubit, b.vertex_type.Z, f"e{b.num_error_bits + 1}")  # Z error bit
    _error(b, qubit, b.vertex_type.X, f"e{b.num_error_bits + 2}")  # X error bit
    b.num_error_bits += 3

def i(b: GraphRepresentation, qubit: int, *_args: float) -> None:
    """Apply identity (advances the row)."""
    ensure_lane(b, qubit)
    v = b.last_vertex[qubit]
    b.graph.set_row(v, last_row(b, qubit) + 1)

def ii(b: GraphRepresentation, qubit1: int, qubit2: int, *_args: float) -> None:
    """Apply two-qubit identity (advances the row on both qubits)."""
    i(b, qubit1)
    i(b, qubit2)

def iswap(b: GraphRepresentation, qubit1: int, qubit2: int) -> None:
    """Swap two qubits and phase the -1 eigenspace of the ZZ observable by i."""
    cnot(b, qubit1, qubit2)
    s(b, qubit2)
    cnot(b, qubit1, qubit2)
    swap(b, qubit1, qubit2)

def iswap_dag(b: GraphRepresentation, qubit1: int, qubit2: int) -> None:
    """Swap two qubits and phase the -1 eigenspace of the ZZ observable by -i."""
    cnot(b, qubit1, qubit2)
    s_dag(b, qubit2)
    cnot(b, qubit1, qubit2)
    swap(b, qubit1, qubit2)

def last_edge(b: GraphRepresentation, qubit: int):
    """Get the last edge for a qubit."""
    edges = b.graph.incident_edges(b.last_vertex[qubit])
    assert len(edges) == 1
    return edges[0]

def last_row(b: GraphRepresentation, qubit: int) -> float:
    """Get the row of the last vertex for a qubit."""
    return b.graph.row(b.last_vertex[qubit])

def m(b: GraphRepresentation, qubit: int, p: float = 0, invert: bool = False) -> None:
    """Measure qubit in Z basis with optional bit-flip error probability p."""
    if invert:
        x(b, qubit)
    _m(b, qubit, p, silent=False)
    if invert:
        x(b, qubit)

def mpad(b: GraphRepresentation, value: int, p: float = 0) -> None:
    """Pad measurement record with a fixed bit value.

    Args:
        b: The graph representation to modify.
        value: The bit value to record (0 or 1).
        p: Error probability for the recorded bit.

    """
    aux = -2
    r(b, aux)
    if value == 1:
        x(b, aux)
    m(b, aux, p=p)

def mpp(
    b: GraphRepresentation,
    paulis: list[tuple[Literal["X", "Y", "Z"], int]],
    invert: bool = False,
    p: float = 0,
) -> None:
    """Measure a single Pauli product.

    Args:
        b: The graph representation to modify.
        paulis: List of (pauli_type, qubit) pairs defining the Pauli product.
        invert: Whether to invert the measurement result.
        p: Measurement flip error probability.

    """
    aux = -2
    r(b, aux)
    h(b, aux)
    _apply_pauli_controls(b, aux, paulis)
    h(b, aux)
    m(b, aux, p=p, invert=invert)

def mr(b: GraphRepresentation, qubit: int, p: float = 0, invert: bool = False) -> None:
    """Z-basis demolition measurement (optionally noisy).

    Projects each target qubit into |0> or |1>, reports its value (false=|0>, true=|1>),
    then resets to |0>.
    """
    m(b, qubit, p=p, invert=invert)
    _r(b, qubit, perform_trace=False)

def mrx(b: GraphRepresentation, qubit: int, p: float = 0, invert: bool = False) -> None:
    """X-basis demolition measurement (optionally noisy).

    Projects each target qubit into |+> or |->, reports its value (false=|+>, true=|->),
    then resets to |+>.
    """
    h(b, qubit)
    m(b, qubit, p=p, invert=invert)
    _r(b, qubit, perform_trace=False)
    h(b, qubit)

def mry(b: GraphRepresentation, qubit: int, p: float = 0, invert: bool = False) -> None:
    """Y-basis demolition measurement (optionally noisy).

    Projects each target qubit into |i> or |-i>, reports its value (false=|i>, true=|-i>),
    then resets to |i>.
    """
    h_yz(b, qubit)
    m(b, qubit, p=p, invert=invert)
    _r(b, qubit, perform_trace=False)
    h_yz(b, qubit)

def mx(b: GraphRepresentation, qubit: int, p: float = 0, invert: bool = False) -> None:
    """Measure qubit in X basis."""
    h(b, qubit)
    m(b, qubit, p=p, invert=invert)
    h(b, qubit)

def mxx(
    b: GraphRepresentation, q0: int, q1: int, p: float = 0, invert: bool = False
) -> None:
    """Measure two qubits in XX basis."""
    mpp(b, [("X", q0), ("X", q1)], invert, p=p)

def my(b: GraphRepresentation, qubit: int, p: float = 0, invert: bool = False) -> None:
    """Measure qubit in Y basis."""
    h_yz(b, qubit)
    m(b, qubit, p=p, invert=invert)
    h_yz(b, qubit)

def myy(
    b: GraphRepresentation, q0: int, q1: int, p: float = 0, invert: bool = False
) -> None:
    """Measure two qubits in YY basis."""
    mpp(b, [("Y", q0), ("Y", q1)], invert, p=p)

def mzz(
    b: GraphRepresentation, q0: int, q1: int, p: float = 0, invert: bool = False
) -> None:
    """Measure two qubits in ZZ basis."""
    mpp(b, [("Z", q0), ("Z", q1)], invert, p=p)

def observable_include(b: GraphRepresentation, rec: list[int], idx: int) -> None:
    """Add observable annotation that XORs the given measurement record bits."""
    idx = int(idx)

    if idx not in b.observables_dict:
        row = min({b.graph.row(b.rec[r]) for r in rec}) - 0.5
        d_rows = {b.graph.row(d) for d in b.detectors + b.observables}
        while row in d_rows:
            row += 1
        v0 = b.graph.add_vertex(VertexType.X, qubit=-1, row=row, phase=f"obs[{idx}]")
        b.observables_dict[idx] = v0

    v0 = b.observables_dict[idx]
    for rec_ in rec:
        b.graph.add_edge((v0, b.rec[rec_]))

def pauli_channel_1(
    b: GraphRepresentation, qubit: int, px: float, py: float, pz: float
) -> None:
    """Apply single-qubit Pauli channel with given X, Y, Z error probabilities."""
    b.channel_probs.append(pauli_channel_1_probs(px, py, pz))
    _error(b, qubit, b.vertex_type.Z, f"e{b.num_error_bits}")
    _error(b, qubit, b.vertex_type.X, f"e{b.num_error_bits + 1}")
    b.num_error_bits += 2

def pauli_channel_2(
    b: GraphRepresentation,
    qubit_i: int,
    qubit_j: int,
    pix: float,
    piy: float,
    piz: float,
    pxi: float,
    pxx: float,
    pxy: float,
    pxz: float,
    pyi: float,
    pyx: float,
    pyy: float,
    pyz: float,
    pzi: float,
    pzx: float,
    pzy: float,
    pzz: float,
) -> None:
    """Apply two-qubit Pauli channel with given error probabilities for all 15 Pauli pairs."""
    b.channel_probs.append(
        pauli_channel_2_probs(
            pix, piy, piz, pxi, pxx, pxy, pxz, pyi, pyx, pyy, pyz, pzi, pzx, pzy, pzz
        )
    )
    _error(b, qubit_i, b.vertex_type.Z, f"e{b.num_error_bits}")
    _error(b, qubit_i, b.vertex_type.X, f"e{b.num_error_bits + 1}")
    _error(b, qubit_j, b.vertex_type.Z, f"e{b.num_error_bits + 2}")
    _error(b, qubit_j, b.vertex_type.X, f"e{b.num_error_bits + 3}")
    b.num_error_bits += 4

def r(b: GraphRepresentation, qubit: int) -> None:
    """Z-basis reset.

    Forces each target qubit into the |0> state by silently measuring it in the Z basis
    and applying an X gate if it ended up in the |1> state.
    """
    _r(b, qubit, perform_trace=True)

def r_x(b: GraphRepresentation, qubit: int, phase: Fraction) -> None:
    """Apply R_X rotation gate with given phase (in units of π)."""
    x_phase(b, qubit, phase)
    b.graph.scalar.add_phase(-phase / 2)

def r_y(b: GraphRepresentation, qubit: int, phase: Fraction) -> None:
    """Apply R_Y rotation gate with given phase (in units of π)."""
    h_yz(b, qubit)
    r_z(b, qubit, phase)
    h_yz(b, qubit)

def r_z(b: GraphRepresentation, qubit: int, phase: Fraction) -> None:
    """Apply R_Z rotation gate with given phase (in units of π)."""
    z_phase(b, qubit, phase)
    b.graph.scalar.add_phase(-phase / 2)

def rx(b: GraphRepresentation, qubit: int) -> None:
    """X-basis reset.

    Forces each target qubit into the |+> state by silently measuring it in the X basis
    and applying a Z gate if it ended up in the |-> state.
    """
    if qubit in b.last_vertex:
        h(b, qubit)
    r(b, qubit)
    h(b, qubit)

def ry(b: GraphRepresentation, qubit: int) -> None:
    """Y-basis reset.

    Forces each target qubit into the |i> state by silently measuring it in the Y basis
    and applying an X gate if it ended up in the |-i> state.
    """
    if qubit in b.last_vertex:
        h_yz(b, qubit)
    r(b, qubit)
    h_yz(b, qubit)

def s(b: GraphRepresentation, qubit: int) -> None:
    """Apply S gate (π/2 Z rotation)."""
    z_phase(b, qubit, Fraction(1, 2))

def s_dag(b: GraphRepresentation, qubit: int) -> None:
    """Apply S† gate (-π/2 Z rotation)."""
    z_phase(b, qubit, Fraction(-1, 2))

def spp(
    b: GraphRepresentation,
    paulis: list[tuple[Literal["X", "Y", "Z"], int]],
    dagger: bool = False,
) -> None:
    """Phase the -1 eigenspace of a Pauli product by i (or -i if dagger).

    Args:
        b: The graph representation to modify.
        paulis: List of (pauli_type, qubit) pairs defining the Pauli product.
        dagger: If True, phase by -i instead of i.

    """
    # Rotate each qubit so its Pauli eigenvalue maps to Z
    for pauli_type, qubit in paulis:
        if pauli_type == "X":
            h(b, qubit)
        elif pauli_type == "Y":
            s_dag(b, qubit)
            h(b, qubit)

    # Accumulate Z-parity into the last qubit
    _, last_qubit = paulis[-1]
    for _, qubit in paulis[:-1]:
        cnot(b, qubit, last_qubit)

    # Phase the parity
    if dagger:
        s_dag(b, last_qubit)
    else:
        s(b, last_qubit)

    # Uncompute parity accumulation
    for _, qubit in reversed(paulis[:-1]):
        cnot(b, qubit, last_qubit)

    # Undo basis rotations
    for pauli_type, qubit in paulis:
        if pauli_type == "X":
            h(b, qubit)
        elif pauli_type == "Y":
            h(b, qubit)
            s(b, qubit)

def sqrt_x(b: GraphRepresentation, qubit: int) -> None:
    """Apply √X gate (π/2 X rotation)."""
    x_phase(b, qubit, Fraction(1, 2))

def sqrt_x_dag(b: GraphRepresentation, qubit: int) -> None:
    """Apply √X† gate (-π/2 X rotation)."""
    x_phase(b, qubit, Fraction(-1, 2))

def sqrt_xx(b: GraphRepresentation, qubit1: int, qubit2: int) -> None:
    """Phases the -1 eigenspace of the XX observable by i."""
    cnot(b, qubit1, qubit2)
    sqrt_x(b, qubit1)
    cnot(b, qubit1, qubit2)

def sqrt_xx_dag(b: GraphRepresentation, qubit1: int, qubit2: int) -> None:
    """Phases the -1 eigenspace of the XX observable by -i."""
    cnot(b, qubit1, qubit2)
    sqrt_x_dag(b, qubit1)
    cnot(b, qubit1, qubit2)

def sqrt_y(b: GraphRepresentation, qubit: int) -> None:
    """Apply √Y gate (π/2 Y rotation)."""
    z(b, qubit)
    h(b, qubit)
    b.graph.scalar.add_phase(Fraction(1, 4))

def sqrt_y_dag(b: GraphRepresentation, qubit: int) -> None:
    """Apply √Y† gate (-π/2 Y rotation)."""
    h(b, qubit)
    z(b, qubit)
    b.graph.scalar.add_phase(Fraction(-1, 4))

def sqrt_yy(b: GraphRepresentation, qubit1: int, qubit2: int) -> None:
    """Phases the -1 eigenspace of the YY observable by i."""
    s(b, qubit1)
    cnot(b, qubit2, qubit1)
    z(b, qubit1)
    h(b, qubit2)
    cnot(b, qubit2, qubit1)
    s(b, qubit1)
    b.graph.scalar.add_phase(Fraction(1, 4))

def sqrt_yy_dag(b: GraphRepresentation, qubit1: int, qubit2: int) -> None:
    """Phases the -1 eigenspace of the YY observable by -i."""
    s_dag(b, qubit1)
    cnot(b, qubit2, qubit1)
    h(b, qubit2)
    z(b, qubit1)
    cnot(b, qubit2, qubit1)
    s_dag(b, qubit1)
    b.graph.scalar.add_phase(Fraction(-1, 4))

def sqrt_z(b: GraphRepresentation, qubit: int) -> None:
    """Apply √Z gate (alias for S gate)."""
    s(b, qubit)

def sqrt_z_dag(b: GraphRepresentation, qubit: int) -> None:
    """Apply √Z† gate (alias for S† gate)."""
    s_dag(b, qubit)