#!/usr/bin/env python3
"""The corpus's exact references: integer arithmetic only, standard library only, sharing no
code with any engine it judges.

Three things live here.

1. **Gate meanings.** `prims()` lowers a gate to what it is by the OpenQASM 3.0 standard
   library's definitions (stdgates.inc), global phase included.

2. **A state vector in big integers.** An amplitude is a pair of integers at scale 2^P, with
   P = PREC bits. Trigonometric constants are computed from an integer pi (Machin's formula)
   and an integer Taylor series with 64 guard bits; an angle that is a multiple of 1/8 turn
   takes its exact value. A rounding per product is at most one unit of 2^-PREC, so after G
   gates on n qubits the state is within about G * 2^(n/2 + 3) * 2^-PREC of the true one in the
   Euclidean norm: for every circuit here, below 2^-180, which is 2^-127 of one unit of a
   double. It is a reference: a difference between an engine and this is the engine's.

3. **A stabilizer tableau in big integers** (Aaronson and Gottesman), and the reduced form of a
   stabilizer group, so that two descriptions of a state compare by equality of bytes.

Distances are computed exactly from the engine's own numbers: a double is a dyadic rational,
and it is converted without rounding before it is compared.
"""
import math
from fractions import Fraction

PREC = 256
GUARD = 64
ONE = 1 << PREC
# A squared weight, at scale 2^(2*PREC), that is rounding and not probability: 2^-448. Every
# outcome in the corpus that can happen has probability far above it.
NOISE = 1 << 64

# ------------------------------------------------------------------ angles


def parse_angle(a):
    """An angle as written in the corpus: "t:M/K" is M/2^K of a full turn, "q:A/B" is A/B of a
    full turn, "r:D" is D radians in decimal. Returns ("t", Fraction of a turn) or
    ("r", Fraction of a radian), both exact."""
    kind, _, v = a.partition(":")
    if kind == "t":
        m, _, k = v.partition("/")
        return ("t", Fraction(int(m), 1 << int(k)))
    if kind == "q":
        num, _, den = v.partition("/")
        return ("t", Fraction(int(num), int(den)))
    if kind == "r":
        return ("r", Fraction(v))
    raise ValueError("an angle is t:M/K, q:A/B or r:D, not %r" % (a,))


def is_dyadic(a):
    return a.startswith("t:") and int(a[2:].partition("/")[2]) <= 32


def pi_scaled(bits):
    """floor(pi * 2^bits) to within a few units, by Machin's formula."""
    g = bits + 32

    def atan_inv(x):
        total = term = (1 << g) // x
        x2 = x * x
        k = 3
        sign = -1
        while term:
            term //= x2
            total += sign * (term // k)
            k += 2
            sign = -sign
        return total

    return (4 * (4 * atan_inv(5) - atan_inv(239))) >> 32


_PI = pi_scaled(PREC + GUARD)


def radians_of(a):
    """The angle in radians as an exact Fraction when written in radians, else to PREC+GUARD
    bits."""
    kind, v = parse_angle(a)
    if kind == "r":
        return v
    return Fraction(v.numerator * 2 * _PI, v.denominator << (PREC + GUARD))


def float_radians(a):
    """The double nearest the angle: Fraction's division rounds correctly, and PREC + GUARD
    bits of pi are far more than a double's rounding can see."""
    return float(radians_of(a))


def _cos_sin_series(x_num, x_den):
    """(cos x, sin x) at scale 2^(PREC+GUARD) for a rational |x| < 1."""
    g = PREC + GUARD
    x = (x_num << g) // x_den
    c = s = 0
    term = 1 << g
    k = 0
    while term:
        if k % 4 == 0:
            c += term
        elif k % 4 == 1:
            s += term
        elif k % 4 == 2:
            c -= term
        else:
            s -= term
        k += 1
        term = (term * x) >> g
        term = term // k if term >= 0 else -((-term) // k)
    return c, s


def cos_sin(a, half=False):
    """(cos, sin) of the angle (half of it with half=True) as integers at scale 2^PREC. A
    multiple of 1/8 turn is exact, up to the rounding of 1/sqrt(2)."""
    kind, v = parse_angle(a)
    if half:
        v = v / 2
    g = PREC + GUARD
    if kind == "t":
        eighths = v * 8
        if eighths.denominator == 1:
            e = eighths.numerator % 8
            r = rsqrt2()
            return [(ONE, 0), (r, r), (0, ONE), (-r, r), (-ONE, 0), (-r, -r), (0, -ONE), (r, -r)][e]
        # Reduce into [-1/8, 1/8) of a turn, by whole quarter turns.
        q = (v * 4 + Fraction(1, 2)).__floor__()
        rest = v - Fraction(q, 4)
        x = Fraction(rest.numerator * 2 * _PI, rest.denominator << g)
    else:
        # In radians: reduce by quarter turns of the integer pi.
        pi = Fraction(_PI, 1 << g)
        q = (v / (pi / 2) + Fraction(1, 2)).__floor__()
        x = v - q * (pi / 2)
    c, s = _cos_sin_series(x.numerator, x.denominator)
    for _ in range(q % 4):
        c, s = -s, c
    return (c >> GUARD, s >> GUARD)


def rsqrt2():
    """1/sqrt(2) at scale 2^PREC."""
    return math.isqrt(1 << (2 * PREC - 1))


# ------------------------------------------------------------------ gate meanings


def _cmul(a, b):
    return ((a[0] * b[0] - a[1] * b[1]) >> PREC, (a[0] * b[1] + a[1] * b[0]) >> PREC)


def _cis(a, half=False):
    return cos_sin(a, half)


def _neg(a):
    kind, _, v = a.partition(":")
    if kind in ("t", "q"):
        m, _, k = v.partition("/")
        return "%s:%d/%s" % (kind, -int(m), k)
    return "r:%s" % (v[1:] if v.startswith("-") else "-" + v)


def _add(a, b):
    """The sum of two angles of one kind, exactly."""
    ka, va = parse_angle(a)
    kb, vb = parse_angle(b)
    if ka != kb:
        raise ValueError("adding angles of two kinds")
    s = va + vb
    if ka == "t":
        if s.denominator & (s.denominator - 1):
            return "q:%d/%d" % (s.numerator, s.denominator)
        return "t:%d/%d" % (s.numerator, s.denominator.bit_length() - 1)
    return "r:%s" % _decimal(s)


def _decimal(f):
    if f.denominator == 1:
        return str(f.numerator)
    d = f.denominator
    k = 0
    while d % 10 == 0:
        d //= 10
        k += 1
    twos = fives = 0
    while d % 2 == 0:
        d //= 2
        twos += 1
    while d % 5 == 0:
        d //= 5
        fives += 1
    if d != 1:
        raise ValueError("not a decimal: %s" % f)
    k += max(twos, fives)
    n = abs(f.numerator) * 10**k // f.denominator
    s = str(n).rjust(k + 1, "0")
    return ("-" if f < 0 else "") + s[:-k] + "." + s[-k:]


_Z = (0, 0)
_1 = (ONE, 0)
_I = (0, ONE)
_MI = (0, -ONE)
_M1 = (-ONE, 0)


def u3_matrix(theta, phi, lam):
    """OpenQASM 3's U(theta, phi, lambda)."""
    c, s = cos_sin(theta, half=True)
    el = _cis(lam)
    ep = _cis(phi)
    epl = _cis(_add(phi, lam))
    return [
        (c, 0),
        ((-el[0] * s) >> PREC, (-el[1] * s) >> PREC),
        ((ep[0] * s) >> PREC, (ep[1] * s) >> PREC),
        ((epl[0] * c) >> PREC, (epl[1] * c) >> PREC),
    ]


def one_qubit_matrix(name, params):
    """A one-qubit gate of the standard library as [m00, m01, m10, m11]."""
    r = rsqrt2()
    half = ONE >> 1
    if name == "id":
        return [_1, _Z, _Z, _1]
    if name == "x":
        return [_Z, _1, _1, _Z]
    if name == "y":
        return [_Z, _MI, _I, _Z]
    if name == "z":
        return [_1, _Z, _Z, _M1]
    if name == "h":
        return [(r, 0), (r, 0), (r, 0), (-r, 0)]
    if name == "s":
        return [_1, _Z, _Z, _I]
    if name == "sdg":
        return [_1, _Z, _Z, _MI]
    if name == "t":
        return [_1, _Z, _Z, (r, r)]
    if name == "tdg":
        return [_1, _Z, _Z, (r, -r)]
    if name == "sx":
        return [(half, half), (half, -half), (half, -half), (half, half)]
    if name == "sxdg":
        return [(half, -half), (half, half), (half, half), (half, -half)]
    if name == "p":
        return [_1, _Z, _Z, _cis(params[0])]
    if name == "rz":
        c, s = cos_sin(params[0], half=True)
        return [(c, -s), _Z, _Z, (c, s)]
    if name == "rx":
        c, s = cos_sin(params[0], half=True)
        return [(c, 0), (0, -s), (0, -s), (c, 0)]
    if name == "ry":
        c, s = cos_sin(params[0], half=True)
        return [(c, 0), (-s, 0), (s, 0), (c, 0)]
    if name == "u":
        return u3_matrix(params[0], params[1], params[2])
    raise ValueError("no one-qubit gate %s" % name)


CONTROLLED = {
    "cx": ("x", 1),
    "cy": ("y", 1),
    "cz": ("z", 1),
    "ch": ("h", 1),
    "cp": ("p", 1),
    "crx": ("rx", 1),
    "cry": ("ry", 1),
    "crz": ("rz", 1),
    "csx": ("sx", 1),
    "ccx": ("x", 2),
    "ccz": ("z", 2),
    "c3x": ("x", 3),
    "c4x": ("x", 4),
    "c3z": ("z", 3),
    "c5z": ("z", 5),
    "c2p": ("p", 2),
    "c2rz": ("rz", 2),
    "c2ry": ("ry", 2),
}


def controlled_form(name):
    """(base gate, number of controls) for a controlled gate's name, or None. `cNx` and `cNz`
    name any number of controls."""
    if name in CONTROLLED:
        return CONTROLLED[name]
    if name.startswith("c") and name[-1] in "xz" and name[1:-1].isdigit():
        return (name[-1], int(name[1:-1]))
    return None


def prims(name, qubits, params):
    """The gate as a list of primitive steps:

      ("m", controls, target, matrix)   a one-qubit matrix under controls
      ("swap", controls, a, b)          exchange two qubits under controls
      ("phase", controls, value)        a phase on the states where the controls are set
    """
    form = controlled_form(name)
    if form is not None:
        base, k = form
        return [("m", list(qubits[:k]), qubits[k], one_qubit_matrix(base, params))]
    if name in ("id", "x", "y", "z", "h", "s", "sdg", "t", "tdg", "sx", "sxdg", "p", "rz", "rx", "ry", "u"):
        return [("m", [], qubits[0], one_qubit_matrix(name, params))]
    if name == "swap":
        return [("swap", [], qubits[0], qubits[1])]
    if name == "cswap":
        return [("swap", [qubits[0]], qubits[1], qubits[2])]
    if name == "cu":
        # stdgates.inc: gate cu(θ, φ, λ, γ) c, t { p(γ) c; ctrl @ U(θ, φ, λ) c, t; }
        theta, phi, lam, gamma = params
        return [("m", [], qubits[0], one_qubit_matrix("p", [gamma])), ("m", [qubits[0]], qubits[1], u3_matrix(theta, phi, lam))]
    if name == "gphase":
        return [("phase", list(qubits), _cis(params[0]))]
    if name == "negcx":
        # negctrl @ x: the target flips where the control is clear.
        return [("m", [], qubits[0], one_qubit_matrix("x", [])), ("m", [qubits[0]], qubits[1], one_qubit_matrix("x", [])), ("m", [], qubits[0], one_qubit_matrix("x", []))]
    if name == "rzz":
        # exp(-i θ/2 Z⊗Z)
        return [("m", [qubits[0]], qubits[1], one_qubit_matrix("x", [])), ("m", [], qubits[1], one_qubit_matrix("rz", params)), ("m", [qubits[0]], qubits[1], one_qubit_matrix("x", []))]
    if name == "rxx":
        h = one_qubit_matrix("h", [])
        return [("m", [], qubits[0], h), ("m", [], qubits[1], h)] + prims("rzz", qubits, params) + [("m", [], qubits[0], h), ("m", [], qubits[1], h)]
    if name == "ryy":
        # rx(π/2): a quarter turn, 1/2^2 of one.
        rx_p = one_qubit_matrix("rx", ["t:1/2"])
        rx_m = one_qubit_matrix("rx", ["t:-1/2"])
        return [("m", [], qubits[0], rx_p), ("m", [], qubits[1], rx_p)] + prims("rzz", qubits, params) + [("m", [], qubits[0], rx_m), ("m", [], qubits[1], rx_m)]
    raise ValueError("no gate %s" % name)


# ------------------------------------------------------------------ the state vector


class State:
    """2^n amplitudes as integers at scale 2^PREC; qubit q is bit q of the index."""

    def __init__(self, n):
        self.n = n
        size = 1 << n
        self.re = [0] * size
        self.im = [0] * size
        self.re[0] = ONE

    def copy(self):
        s = State.__new__(State)
        s.n, s.re, s.im = self.n, list(self.re), list(self.im)
        return s

    def apply_prim(self, p):
        if p[0] == "m":
            self._matrix(p[1], p[2], p[3])
        elif p[0] == "swap":
            self._swap(p[1], p[2], p[3])
        else:
            self._phase(p[1], p[2])

    def _matrix(self, controls, target, m):
        re, im = self.re, self.im
        bit = 1 << target
        cmask = 0
        for c in controls:
            cmask |= 1 << c
        (a_r, a_i), (b_r, b_i), (c_r, c_i), (d_r, d_i) = m
        p = PREC
        size = 1 << self.n
        diag = b_r == 0 and b_i == 0 and c_r == 0 and c_i == 0
        trivial_a = diag and a_r == ONE and a_i == 0
        for i in range(size):
            if i & bit or (i & cmask) != cmask:
                continue
            j = i | bit
            xr, xi, yr, yi = re[i], im[i], re[j], im[j]
            if diag:
                if not trivial_a:
                    re[i] = (a_r * xr - a_i * xi) >> p
                    im[i] = (a_r * xi + a_i * xr) >> p
                re[j] = (d_r * yr - d_i * yi) >> p
                im[j] = (d_r * yi + d_i * yr) >> p
            else:
                re[i] = (a_r * xr - a_i * xi + b_r * yr - b_i * yi) >> p
                im[i] = (a_r * xi + a_i * xr + b_r * yi + b_i * yr) >> p
                re[j] = (c_r * xr - c_i * xi + d_r * yr - d_i * yi) >> p
                im[j] = (c_r * xi + c_i * xr + d_r * yi + d_i * yr) >> p

    def _swap(self, controls, a, b):
        re, im = self.re, self.im
        ab, bb = 1 << a, 1 << b
        cmask = 0
        for c in controls:
            cmask |= 1 << c
        for i in range(1 << self.n):
            if (i & ab) and not (i & bb) and (i & cmask) == cmask:
                j = (i ^ ab) | bb
                re[i], re[j] = re[j], re[i]
                im[i], im[j] = im[j], im[i]

    def _phase(self, controls, value):
        re, im = self.re, self.im
        cmask = 0
        for c in controls:
            cmask |= 1 << c
        vr, vi = value
        for i in range(1 << self.n):
            if (i & cmask) == cmask:
                re[i], im[i] = (vr * re[i] - vi * im[i]) >> PREC, (vr * im[i] + vi * re[i]) >> PREC

    def gate(self, name, qubits, params):
        for p in prims(name, qubits, params):
            self.apply_prim(p)

    def prob_one(self, q):
        """The weight on states where q reads 1, at scale 2^(2*PREC)."""
        bit = 1 << q
        return sum(self.re[i] * self.re[i] + self.im[i] * self.im[i] for i in range(1 << self.n) if i & bit)

    def norm2(self):
        return sum(x * x for x in self.re) + sum(x * x for x in self.im)

    def collapse(self, q, outcome):
        """Keep the branch where q reads outcome, unnormalised: its squared norm is the
        probability of the path that reached it."""
        bit = 1 << q
        for i in range(1 << self.n):
            if bool(i & bit) != bool(outcome):
                self.re[i] = 0
                self.im[i] = 0

    def flip(self, q):
        bit = 1 << q
        for i in range(1 << self.n):
            if not i & bit:
                j = i | bit
                self.re[i], self.re[j] = self.re[j], self.re[i]
                self.im[i], self.im[j] = self.im[j], self.im[i]


def run_unitary(n, steps):
    """The state after a program with no measurement."""
    s = State(n)
    for g in steps:
        s.gate(g["g"], g["q"], g.get("a", []))
    return s


def distribution(n, n_bits, steps):
    """The exact distribution of a program's classical bits, by walking every branch of its
    measurements and resets: {bit string with bit 0 last: Fraction}. A step is a gate, or
    {"measure": q, "bit": c}, {"reset": q}, or {"if": [c, v], "then": [gates]} on one bit."""

    # Trailing measurements read off the last state, with no branching.
    k = len(steps)
    while k > 0 and "measure" in steps[k - 1]:
        k -= 1
    tail = steps[k:]
    head = steps[:k]
    tail_qubits = [t["measure"] for t in tail]
    if len(set(tail_qubits)) != len(tail_qubits):
        head, tail = steps, []
    out = {}

    def finish(state, bits):
        if not tail:
            w = state.norm2()
            if w > NOISE:
                key = "".join("1" if bits[c] else "0" for c in range(n_bits - 1, -1, -1))
                out[key] = out.get(key, 0) + w
            return
        for i in range(1 << n):
            w = state.re[i] * state.re[i] + state.im[i] * state.im[i]
            if w <= NOISE:
                continue
            b = list(bits)
            for t in tail:
                b[t["bit"]] = (i >> t["measure"]) & 1
            key = "".join("1" if b[c] else "0" for c in range(n_bits - 1, -1, -1))
            out[key] = out.get(key, 0) + w

    def walk(state, bits, at):
        while at < len(head):
            st = head[at]
            at += 1
            if "g" in st:
                state.gate(st["g"], st["q"], st.get("a", []))
            elif "if" in st:
                c, v = st["if"]
                if bits[c] == v:
                    for g in st["then"]:
                        state.gate(g["g"], g["q"], g.get("a", []))
            elif "measure" in st or "reset" in st:
                q = st.get("measure", st.get("reset"))
                total = state.norm2()
                w1 = state.prob_one(q)
                w0 = total - w1
                # A branch whose weight is the arithmetic's own rounding is no branch.
                branches = []
                if w0 > NOISE:
                    branches.append(0)
                if w1 > NOISE:
                    branches.append(1)
                for idx, o in enumerate(branches):
                    s2 = state if idx == len(branches) - 1 else state.copy()
                    s2.collapse(q, o)
                    b2 = list(bits)
                    if "measure" in st:
                        b2[st["bit"]] = o
                    elif o == 1:
                        s2.flip(q)
                    walk(s2, b2, at)
                return
            else:
                raise ValueError("no step %r" % (st,))
        finish(state, bits)

    walk(State(n), [0] * n_bits, 0)
    total = sum(out.values())
    return {k: Fraction(v, total) for k, v in sorted(out.items())}


# ------------------------------------------------------------------ comparison


def ints_from_floats(xs):
    """Doubles to integers at scale 2^PREC. A double is dyadic; one below 2^-PREC of a unit
    rounds, and that is far below anything compared here."""
    out = []
    for x in xs:
        a, b = x.as_integer_ratio()
        out.append((a << PREC) // b)
    return out


def ints_from_fixed(xs, frac=30):
    return [x << (PREC - frac) for x in xs]


def distance(ref_re, ref_im, re, im):
    """(||a - b||, ||a - e^{iφ} b|| minimised over φ, the phase φ in radians, ||a||^2, ||b||^2),
    as floats computed exactly from the integers."""
    s = 2 * PREC
    d2 = 0
    ip_r = 0
    ip_i = 0
    na = 0
    nb = 0
    for ar, ai, br, bi in zip(ref_re, ref_im, re, im):
        dr = ar - br
        di = ai - bi
        d2 += dr * dr + di * di
        # <b|a>: the phase that takes b to a.
        ip_r += br * ar + bi * ai
        ip_i += br * ai - bi * ar
        na += ar * ar + ai * ai
        nb += br * br + bi * bi
    mod = math.isqrt(ip_r * ip_r + ip_i * ip_i)
    dm2 = max(na + nb - 2 * mod, 0)
    fl = lambda v, shift: float(Fraction(v, 1 << shift))
    return (
        math.sqrt(fl(d2, s)),
        math.sqrt(fl(dm2, s)),
        math.atan2(fl(ip_i, s), fl(ip_r, s)),
        fl(na, s),
        fl(nb, s),
    )


def amplitude_error(ref, re, im):
    """|a - b| for one amplitude: ref at scale 2^PREC, re and im doubles."""
    r = ints_from_floats([re, im])
    dr, di = ref[0] - r[0], ref[1] - r[1]
    return math.sqrt(float(Fraction(dr * dr + di * di, 1 << (2 * PREC))))


# ------------------------------------------------------------------ closed forms


def qft_amplitude(n, x, y):
    """<y| QFT |x>, with qubit 0 the lowest bit of the index: e^{2πi x y / 2^n} / sqrt(2^n)."""
    t = "t:%d/%d" % ((x * y) % (1 << n), n)
    c, s = cos_sin(t)
    root = math.isqrt(1 << (2 * PREC - n))
    return ((c * root) >> PREC, (s * root) >> PREC)


def qpe_third_law(bits, y):
    """P(y) for phase estimation of a phase of 1/3 on `bits` counting qubits, at scale
    2^(2·PREC): sin²(π·2^t·δ) / (2^(2t)·sin²(π·δ)) with δ = 1/3 − y/2^t."""
    d = Fraction(1, 3) - Fraction(y, 1 << bits)
    top = d * (1 << bits)
    s_top = cos_sin("q:%d/%d" % (top.numerator, top.denominator), half=True)[1]
    s_bot = cos_sin("q:%d/%d" % (d.numerator, d.denominator), half=True)[1]
    return (s_top * s_top << (2 * PREC)) // (s_bot * s_bot << (2 * bits))


def grover_law(n, k):
    """sin²((2k+1)·θ) with sin θ = 2^(−n/2), at scale 2^(2·PREC): the rotation by 2θ applied k
    times to (cos θ, sin θ), in integers of PREC + GUARD bits."""
    g = PREC + GUARD
    s = math.isqrt(1 << (2 * g - n))
    c = math.isqrt((1 << (2 * g)) - (1 << (2 * g - n)))
    c2, s2 = (c * c - s * s) >> g, (2 * c * s) >> g
    x, y = c, s
    for _ in range(k):
        x, y = (x * c2 - y * s2) >> g, (y * c2 + x * s2) >> g
    return (y * y) >> (2 * GUARD)


def _inner(a_re, a_im, b_re, b_im):
    """<a|b> at scale 2^(2·PREC)."""
    r = i = 0
    for ar, ai, br, bi in zip(a_re, a_im, b_re, b_im):
        r += ar * br + ai * bi
        i += ar * bi - ai * br
    return r, i


def _turned(v_re, v_im, ip):
    """v times ip/|ip|: the global phase that brings v nearest the state ip was taken against."""
    m = math.isqrt(ip[0] * ip[0] + ip[1] * ip[1])
    if m == 0:
        return list(v_re), list(v_im)
    ur, ui = (ip[0] << PREC) // m, (ip[1] << PREC) // m
    return [(r * ur - i * ui) >> PREC for r, i in zip(v_re, v_im)], [(r * ui + i * ur) >> PREC for r, i in zip(v_re, v_im)]


def held_share(ideal, without, v_re, v_im):
    """How much of the step from `without` (the circuit with one gate left out) to `ideal` an
    answer took: Re<Δ, v − w> / |Δ|², Δ = ideal − w, each state first turned by the global phase
    that brings it nearest the ideal. 1 holds the gate and 0 drops it. In integers at scale
    2^PREC, so a step of 2^-60, which a double cannot resolve beside an amplitude of 1, is seen."""
    w_re, w_im = _turned(without.re, without.im, _inner(without.re, without.im, ideal.re, ideal.im))
    a_re, a_im = _turned(v_re, v_im, _inner(v_re, v_im, ideal.re, ideal.im))
    num = den = 0
    for ir, ii, wr, wi, ar, ai in zip(ideal.re, ideal.im, w_re, w_im, a_re, a_im):
        dr, di = ir - wr, ii - wi
        num += dr * (ar - wr) + di * (ai - wi)
        den += dr * dr + di * di
    return float(Fraction(num, den)) if den else float("nan")


# ------------------------------------------------------------------ the stabilizer tableau


class Tableau:
    """Aaronson and Gottesman's tableau, stored by qubit: x[q] and z[q] are integers whose bit r
    is row r (rows 0..n-1 destabilizers, n..2n-1 stabilizers), and r holds the signs."""

    def __init__(self, n):
        self.n = n
        self.x = [1 << q for q in range(n)]
        self.z = [1 << (q + n) for q in range(n)]
        self.r = 0

    def h(self, q):
        x, z = self.x[q], self.z[q]
        self.r ^= x & z
        self.x[q], self.z[q] = z, x

    def s(self, q):
        x, z = self.x[q], self.z[q]
        self.r ^= x & z
        self.z[q] = z ^ x

    def sdg(self, q):
        self.z_gate(q)
        self.s(q)

    def x_gate(self, q):
        self.r ^= self.z[q]

    def z_gate(self, q):
        self.r ^= self.x[q]

    def y_gate(self, q):
        self.r ^= self.x[q] ^ self.z[q]

    def cx(self, a, b):
        xa, za, xb, zb = self.x[a], self.z[a], self.x[b], self.z[b]
        rows = (1 << (2 * self.n)) - 1
        self.r ^= xa & zb & ((xb ^ za) ^ rows)
        self.x[b] = xb ^ xa
        self.z[a] = za ^ zb

    def cz(self, a, b):
        self.h(b)
        self.cx(a, b)
        self.h(b)

    def swap(self, a, b):
        self.cx(a, b)
        self.cx(b, a)
        self.cx(a, b)

    def sx(self, q):
        self.h(q)
        self.s(q)
        self.h(q)

    def sxdg(self, q):
        self.h(q)
        self.sdg(q)
        self.h(q)

    def cy(self, a, b):
        self.sdg(b)
        self.cx(a, b)
        self.s(b)

    def gate(self, name, q):
        fn = {"h": self.h, "s": self.s, "sdg": self.sdg, "x": self.x_gate, "y": self.y_gate, "z": self.z_gate, "sx": self.sx, "sxdg": self.sxdg, "cx": self.cx, "cz": self.cz, "cy": self.cy, "swap": self.swap, "id": lambda _: None}[name]
        fn(*q)

    def stabilizers(self):
        """The stabilizers as (x, z, sign) with x and z integers whose bit q is qubit q."""
        n = self.n
        xs = [0] * n
        zs = [0] * n
        # A transpose that visits only the set bits, so a sparse tableau costs what it holds.
        for q in range(n):
            for cols, out in ((self.x[q] >> n, xs), (self.z[q] >> n, zs)):
                v = cols
                while v:
                    low = v & -v
                    out[low.bit_length() - 1] |= 1 << q
                    v ^= low
        return [(xs[i], zs[i], (self.r >> (n + i)) & 1) for i in range(n)]


def pauli_mul(a, b):
    """The product of two commuting Hermitian Paulis (x, z, sign), sign included, where a row
    (x, z, s) means (-1)^s i^(x·z) X^x Z^z, which is the string with Y where x and z meet."""
    x1, z1, s1 = a
    x2, z2, s2 = b
    x3, z3 = x1 ^ x2, z1 ^ z2
    e = (x1 & z1).bit_count() + (x2 & z2).bit_count() - (x3 & z3).bit_count() + 2 * (z1 & x2).bit_count() + 2 * (s1 + s2)
    e %= 4
    if e % 2:
        raise ValueError("the product of two commuting Paulis is Hermitian")
    return (x3, z3, e // 2)


def reduced_group(rows, n):
    """The reduced row echelon form of a stabilizer group over (x bits, then z bits), qubit 0
    first, with each row's sign carried through its products. The same group gives the same
    rows, whatever generators it arrived as. Also returns the rank and the rank of the x part."""
    rows = list(rows)
    out = []
    x_rank = 0
    for column in ("x", "z"):
        for q in range(n):
            bit = 1 << q
            at = 0 if column == "x" else 1
            pivot = None
            for i, r in enumerate(rows):
                if r[at] & bit:
                    pivot = i
                    break
            if pivot is None:
                continue
            p = rows.pop(pivot)
            rows = [pauli_mul(r, p) if r[at] & bit else r for r in rows]
            out = [pauli_mul(r, p) if r[at] & bit else r for r in out]
            out.append(p)
            if column == "x":
                x_rank += 1
    return out, len(out), x_rank


def pauli_text(row, n):
    """A row as its sign and XYZ_ letters, qubit 0 first."""
    x, z, s = row
    letters = []
    for q in range(n):
        xb, zb = (x >> q) & 1, (z >> q) & 1
        letters.append("_XZY"[xb + 2 * zb])
    return ("-" if s else "+") + "".join(letters)


def row_from_text(text):
    """A signed string of letters XYZ_ (or I), qubit 0 first, as a row."""
    s = 1 if text[0] == "-" else 0
    body = text[1:] if text[0] in "+-" else text
    x = z = 0
    for q, c in enumerate(body):
        if c == "X":
            x |= 1 << q
        elif c == "Z":
            z |= 1 << q
        elif c == "Y":
            x |= 1 << q
            z |= 1 << q
    return (x, z, s)


def commute(a, b):
    return ((a[0] & b[1]).bit_count() + (a[1] & b[0]).bit_count()) % 2 == 0


# ------------------------------------------------------------------ Pauli strings


def pauli_masks(text, n):
    """A Pauli string written as the engine takes it ("X1 Z2", qubit after the letter, an optional
    leading coefficient) as (coefficient as a Fraction, x mask, z mask, number of Ys)."""
    coeff = Fraction(1)
    x = z = ny = 0
    for tok in text.split():
        if tok[0] in "XYZ" and tok[1:].isdigit():
            q = int(tok[1:])
            if q >= n:
                raise ValueError("qubit %d of %r past %d" % (q, text, n))
            if tok[0] in "XY":
                x |= 1 << q
            if tok[0] in "ZY":
                z |= 1 << q
            ny += tok[0] == "Y"
        else:
            coeff *= Fraction(tok)
    return coeff, x, z, ny


def _pauli_phase(k, z, ny):
    """P|k> = phase(k) |k ^ x>, P = i^ny X^x Z^z, as a power of i (0..3)."""
    return (ny + 2 * ((k & z).bit_count() & 1)) % 4


def _times_i_power(re, im, p):
    return [(re, im), (-im, re), (-re, -im), (im, -re)][p]


def pure_expectation(state, text):
    """<psi|P|psi> of a State, exactly, as a Fraction (the string's coefficient included)."""
    coeff, x, z, ny = pauli_masks(text, state.n)
    tr = 0
    for k in range(1 << state.n):
        a_r, a_i = state.re[k], state.im[k]
        if not (a_r or a_i):
            continue
        pr, pi = _times_i_power(a_r, a_i, _pauli_phase(k, z, ny))
        b_r, b_i = state.re[k ^ x], state.im[k ^ x]
        # conj(b) * (phase * a), real part: the imaginary parts cancel over the sum.
        tr += b_r * pr + b_i * pi
    return coeff * Fraction(tr, 1 << (2 * PREC))


# ------------------------------------------------------------------ the density matrix


def _conj_matrix(m):
    return [(r, -i) for r, i in m]


def _dagger_matrix(m):
    (a, b, c, d) = m
    return [(a[0], -a[1]), (c[0], -c[1]), (b[0], -b[1]), (d[0], -d[1])]


def dagger_prims(name, qubits, params):
    """The gate's inverse as primitive steps: each step of prims() undone, in reverse order."""
    out = []
    for p in reversed(prims(name, qubits, params)):
        if p[0] == "m":
            out.append(("m", p[1], p[2], _dagger_matrix(p[3])))
        elif p[0] == "swap":
            out.append(p)
        else:
            out.append(("phase", p[1], (p[2][0], -p[2][1])))
    return out


_PAULI_1 = {"X": [_Z, _1, _1, _Z], "Y": [_Z, _MI, _I, _Z], "Z": [_1, _Z, _Z, _M1]}


class Density:
    """A density matrix of n qubits as 4^n integer pairs at scale 2^PREC: entry (i, j) at index
    i | j << n, so row qubit q is bit q and column qubit q is bit q + n. A gate U acts as U on the
    row bits and its conjugate on the column bits (vec(UρU†) = (U ⊗ Ū) vec ρ), and a channel as a
    sum of such terms. Every step rounds each entry by at most a unit or two of 2^-PREC."""

    def __init__(self, n):
        self.n = n
        self.s = State(2 * n)

    def apply_prims(self, ps):
        n, s = self.n, self.s
        for p in ps:
            if p[0] == "m":
                s._matrix(p[1], p[2], p[3])
                s._matrix([c + n for c in p[1]], p[2] + n, _conj_matrix(p[3]))
            elif p[0] == "swap":
                s._swap(p[1], p[2], p[3])
                s._swap([c + n for c in p[1]], p[2] + n, p[3] + n)
            else:
                s._phase(p[1], p[2])
                s._phase([c + n for c in p[1]], (p[2][0], -p[2][1]))

    def mix(self, terms):
        """ρ ← Σ w · K ρ K†, for terms [(w, [(qubit, 2x2 matrix), ...])], w a Fraction."""
        n = self.n
        size = len(self.s.re)
        re, im = [0] * size, [0] * size
        for w, ops in terms:
            if w == 0:
                continue
            c = self.s.copy()
            for q, m in ops:
                c._matrix([], q, m)
                c._matrix([], q + n, _conj_matrix(m))
            num, den = w.numerator, w.denominator
            for i in range(size):
                if c.re[i]:
                    re[i] += (c.re[i] * num) // den
                if c.im[i]:
                    im[i] += (c.im[i] * num) // den
        self.s.re, self.s.im = re, im

    def depolarize1(self, q, p):
        self.mix([(1 - p, [])] + [(p / 3, [(q, _PAULI_1[a])]) for a in "XYZ"])

    def depolarize2(self, a, b, p):
        terms = [(1 - p, [])]
        for j in range(1, 16):
            ops = []
            for q, letter in ((a, "IXYZ"[j >> 2]), (b, "IXYZ"[j & 3])):
                if letter != "I":
                    ops.append((q, _PAULI_1[letter]))
            terms.append((p / 15, ops))
        self.mix(terms)

    def dephase(self, q, p):
        self.mix([(1 - p, []), (p, [(q, _PAULI_1["Z"])])])

    def damp(self, q, g):
        """Amplitude damping of strength γ: K0 = diag(1, √(1−γ)), K1 = √γ |0><1|, the roots to
        the nearest unit below."""
        r1 = math.isqrt(((g.denominator - g.numerator) << (2 * PREC)) // g.denominator)
        rg = math.isqrt((g.numerator << (2 * PREC)) // g.denominator)
        self.mix([(Fraction(1), [(q, [_1, _Z, _Z, (r1, 0)])]), (Fraction(1), [(q, [_Z, (rg, 0), _Z, _Z])])])

    def diagonal(self):
        """The probabilities of the computational basis states, at scale 2^PREC."""
        n = self.n
        return [self.s.re[i | i << n] for i in range(1 << n)]

    def expectation(self, text):
        """Tr(Pρ) as a Fraction, the string's coefficient included."""
        n = self.n
        coeff, x, z, ny = pauli_masks(text, n)
        tr = 0
        for k in range(1 << n):
            # Tr(Pρ) = Σ_k phase(k) ρ[k, k ^ x]
            idx = k | ((k ^ x) << n)
            r, i = _times_i_power(self.s.re[idx], self.s.im[idx], _pauli_phase(k, z, ny))
            tr += r
        return coeff * Fraction(tr, ONE)

    def purity(self):
        """Tr(ρ²) = Σ |ρ_ij|², as a Fraction."""
        return Fraction(sum(r * r + i * i for r, i in zip(self.s.re, self.s.im)), 1 << (2 * PREC))


NOISE_KEYS = ("depolarizing_1q", "depolarizing_2q", "dephasing", "amplitude_damping", "readout")


def noise_of(noise):
    """A job's noise as exact Fractions, from the decimal text of each probability."""
    out = {k: Fraction(0) for k in NOISE_KEYS}
    for k, v in (noise or {}).items():
        if k not in out:
            raise ValueError("no noise %s" % k)
        out[k] = Fraction(repr(v)) if isinstance(v, float) else Fraction(v)
    return out


def noisy_density(n, gates, noise, fold=0):
    """The exact noisy state of a program of gates as the engine's noise model places its channels
    (wai-quantum quantum_noisy::DensityMatrix::evolve): after each gate, a two-qubit gate's
    depolarizing channel on its pair; then for each qubit it touched, the one-qubit depolarizing
    channel (a gate on one qubit or on three or more), dephasing and damping. `fold` runs U (U†U)^k,
    each gate of U† the inverse of a gate of U as one gate, with its channels after it."""
    nz = noise_of(noise)
    rho = Density(n)
    order = [(g, False) for g in gates]
    back = [(g, True) for g in reversed(gates)]
    for _ in range(fold):
        order += back + [(g, False) for g in gates]
    for g, inverse in order:
        rho.apply_prims(dagger_prims(g["g"], g["q"], g.get("a", [])) if inverse else prims(g["g"], g["q"], g.get("a", [])))
        qs = g["q"]
        k = len(qs)
        if k == 2 and nz["depolarizing_2q"]:
            rho.depolarize2(qs[0], qs[1], nz["depolarizing_2q"])
        for q in qs:
            if k != 2 and nz["depolarizing_1q"]:
                rho.depolarize1(q, nz["depolarizing_1q"])
            if nz["dephasing"]:
                rho.dephase(q, nz["dephasing"])
            if nz["amplitude_damping"]:
                rho.damp(q, nz["amplitude_damping"])
    return rho


def noisy_distribution(rho, measured, n_bits, readout=Fraction(0)):
    """The exact law of the classical bits from a density matrix's diagonal, each measured bit then
    misread with probability `readout`: {bits, bit 0 last: Fraction}. `measured` is [(qubit, bit)]."""
    out = {}
    for x, w in enumerate(rho.diagonal()):
        if w <= 0:
            continue
        bits = [0] * n_bits
        for q, b in measured:
            bits[b] = (x >> q) & 1
        key = "".join(str(bits[c]) for c in range(n_bits - 1, -1, -1))
        out[key] = out.get(key, 0) + w
    total = sum(out.values())
    law = {k: Fraction(v, total) for k, v in out.items()}
    if readout:
        for b in range(n_bits):
            pos = n_bits - 1 - b
            new = {}
            for k, p in law.items():
                flip = k[:pos] + ("1" if k[pos] == "0" else "0") + k[pos + 1:]
                new[k] = new.get(k, 0) + p * (1 - readout)
                new[flip] = new.get(flip, 0) + p * readout
            law = new
    return {k: v for k, v in sorted(law.items()) if v}


# ------------------------------------------------------------------ closed forms of a driven qubit


def cos_sin_turns(v):
    """(cos, sin) of an angle of v turns (a Fraction), at scale 2^PREC: cos_sin for a value that
    is no written angle."""
    g = PREC + GUARD
    q = (v * 4 + Fraction(1, 2)).__floor__()
    rest = v - Fraction(q, 4)
    x = Fraction(rest.numerator * 2 * _PI, rest.denominator << g)
    c, s = _cos_sin_series(x.numerator, x.denominator)
    for _ in range(q % 4):
        c, s = -s, c
    return (c >> GUARD, s >> GUARD)


def sqrt_fraction(x, bits=PREC + GUARD):
    """√x for a non-negative Fraction, to 2^-bits."""
    return Fraction(math.isqrt((x.numerator << (2 * bits)) // x.denominator), 1 << bits)


def rabi_p1(rabi_hz, amplitude, seconds, detuning_hz=Fraction(0)):
    """P(1) after a constant drive from |0> in the rotating frame (the rotating-wave model): the
    Rabi formula r²/g² · sin²(π g t), r = rabi·|A| and g = √(r² + δ²) in cycles a second. Exact
    up to 2^-PREC."""
    r = rabi_hz * amplitude
    g2 = r * r + detuning_hz * detuning_hz
    if g2 == 0:
        return Fraction(0)
    turns = sqrt_fraction(g2 * seconds * seconds / 4)
    _, s = cos_sin_turns(turns)
    return (r * r / g2) * Fraction(s * s, 1 << (2 * PREC))


def density_entries(rho):
    """The density matrix as a dict {(i, j): complex Fraction pair (re, im)}, for small n."""
    n = rho.n
    out = {}
    for i in range(1 << n):
        for j in range(1 << n):
            idx = i | j << n
            out[(i, j)] = (Fraction(rho.s.re[idx], ONE), Fraction(rho.s.im[idx], ONE))
    return out


def trace_rho2_pauli(rho, text):
    """Tr(ρ² P) exactly, from the integer density matrix (small n): Σ_k phase(k)·(ρ²)[k, k ^ x]."""
    n = rho.n
    coeff, x, z, ny = pauli_masks(text, n)
    m = density_entries(rho)
    d = 1 << n
    tr = Fraction(0)
    for k in range(d):
        r, i = Fraction(0), Fraction(0)
        for l in range(d):
            ar, ai = m[(k, l)]
            br, bi = m[(l, k ^ x)]
            r += ar * br - ai * bi
            i += ar * bi + ai * br
        pr, _ = _times_i_power(r, i, _pauli_phase(k, z, ny))
        tr += pr
    return coeff * tr