"""Board emitter-graph: deterministic photonic graph-state generation from a few quantum emitters (SPEC_v2_final s.5).

Target: the graph state on photons 0..n-1 with stabilizers K_v = X_v prod_{u ~ v} Z_u, all signs +1.
Artifact {"emitters": E, "ops": [...], "local_cliffords": [n], "photon_corrections": [n]}.
Ops act on emitters 0..E-1 only, except ["EMIT", e, v]: fresh photon v in |0>, then CX(e -> v).
Tags RZ RX (prepare |0>, |+>; only on an uninitialized or measured emitter), H S S_DAG X Z, CX CZ (two distinct
emitters), MZ MX (destructive; one record bit, 0 = +1 eigenstate), PX PZ [e, bits, value] (apply iff the XOR of the
listed earlier records equals value, the ft-stateprep convention). After the ops, photon v gets its local Clifford
U_v (U X U^dag = image x, U Z U^dag = image z), then X^a Z^b with a = x_const + sum of records in x, b likewise (mod 2).
Accept iff every K_v (identity on the emitters) lies in the final stabilizer group with sign +1 for every outcome;
the n generators fix the photon state uniquely, so this is "the photons are exactly the target and unentangled
with the emitters, in every branch". Score (ee_gates, emitters) = (number of CX and CZ, E), both minimized.

Exact method: an Aaronson-Gottesman stabilizer/destabilizer tableau over Q = n + E qubits (photons first), stored by
columns as int bitmasks over rows, so a gate is O(1) big-int operations. Stabilizer signs are affine functions over
GF(2) of the independent random measurement outcomes: a constant bit (bitmask r0 over rows) plus a variable part per
row (int, bit j + 1 = j-th random outcome). A measurement is either random (new variable) or deterministic (record =
the affine sign of the product of stabilizers equal to the measured Pauli). After a measurement the qubit is
decoupled (its stabilizer row is +-Z_q, its destabilizer X_q, no other row touches it), so a later preparation just
overwrites that row. No branch is enumerated or sampled.
"""
from common import Reject, exact_keys, want_int, want_list, is_int, load_instance, short

popcount = int.bit_count if hasattr(int, "bit_count") else (lambda v: bin(v).count("1"))

ONE = ("RZ", "RX", "H", "S", "S_DAG", "X", "Z", "MZ", "MX")
PAULI = {"X": (1, 0), "Y": (1, 1), "Z": (0, 1)}
IMAGES = ("+X", "-X", "+Y", "-Y", "+Z", "-Z")


def bits(m):
    while m:
        low = m & -m
        yield low.bit_length() - 1
        m ^= low


def prefix_parity(a, width):
    """Bit l of the result = XOR of the bits of a strictly below l (for l < width)."""
    t, s = a << 1, 1
    while s < width:
        t ^= t << s
        s <<= 1
    return t & ((1 << width) - 1)


class Tableau:
    def __init__(self, Q):
        self.Q = Q
        self.sx, self.sz = [0] * Q, [1 << q for q in range(Q)]   # stabilizer row i = +Z_i
        self.dx, self.dz = [1 << q for q in range(Q)], [0] * Q   # destabilizer row i = X_i (signs unused)
        self.r0 = 0                                              # constant sign bits of the stabilizer rows
        self.rv = [0] * Q                                        # variable sign parts (bit 0 always clear)
        self.owner = list(range(Q))                              # decoupled qubit q -> its stabilizer row
        self.nvar = 0

    # Clifford gates (Aaronson-Gottesman update rules applied to whole columns)
    def h(self, q):
        self.r0 ^= self.sx[q] & self.sz[q]
        self.sx[q], self.sz[q] = self.sz[q], self.sx[q]
        self.dx[q], self.dz[q] = self.dz[q], self.dx[q]

    def s(self, q):
        self.r0 ^= self.sx[q] & self.sz[q]
        self.sz[q] ^= self.sx[q]
        self.dz[q] ^= self.dx[q]

    def x(self, q):
        self.r0 ^= self.sz[q]

    def z(self, q):
        self.r0 ^= self.sx[q]

    def cx(self, a, b):
        sx, sz = self.sx, self.sz
        self.r0 ^= sx[a] & sz[b] & ~(sx[b] ^ sz[a])
        sx[b] ^= sx[a]
        sz[a] ^= sz[b]
        self.dx[b] ^= self.dx[a]
        self.dz[a] ^= self.dz[b]

    def cz(self, a, b):
        sx, sz = self.sx, self.sz
        self.r0 ^= sx[a] & sx[b] & (sz[a] ^ sz[b])
        sz[a] ^= sx[b]
        sz[b] ^= sx[a]
        self.dz[a] ^= self.dx[b]
        self.dz[b] ^= self.dx[a]

    def clifford(self, q, img):
        """Single-qubit Clifford given by the signed images of X and Z (strings such as '-Y')."""
        px, pz = img["x"][1], img["z"][1]
        py = ({"X", "Y", "Z"} - {px, pz}).pop()
        neg_y = (img["x"][0] == "-") ^ (img["z"][0] == "-") ^ ((px, pz) in (("Z", "X"), ("X", "Y"), ("Y", "Z")))
        image = {"X": (px, img["x"][0] == "-"), "Z": (pz, img["z"][0] == "-"), "Y": (py, neg_y)}
        sx, sz, dx, dz = self.sx[q], self.sz[q], self.dx[q], self.dz[q]
        masks = {"X": (sx & ~sz, dx & ~dz), "Y": (sx & sz, dx & dz), "Z": (sz & ~sx, dz & ~dx)}
        nsx = nsz = ndx = ndz = flip = 0
        for t, (ms, md) in masks.items():
            p, neg = image[t]
            if PAULI[p][0]:
                nsx, ndx = nsx | ms, ndx | md
            if PAULI[p][1]:
                nsz, ndz = nsz | ms, ndz | md
            if neg:
                flip |= ms
        self.sx[q], self.sz[q], self.dx[q], self.dz[q] = nsx, nsz, ndx, ndz
        self.r0 ^= flip

    def flip(self, mask, aff):
        """XOR the affine expression aff (bit 0 constant, bit j + 1 variable j) into the signs of rows in mask."""
        if aff & 1:
            self.r0 ^= mask
        var = aff & ~1
        if var:
            for i in bits(mask):
                self.rv[i] ^= var

    def row(self, cols_x, cols_z, i):
        b = 1 << i
        return [k for k in range(self.Q) if cols_x[k] & b], [k for k in range(self.Q) if cols_z[k] & b]

    def product(self, J):
        """Product of the stabilizer rows in mask J (increasing order): (x bits, z bits, affine sign).

        A row is the Hermitian Pauli (-1)^r i^y X^x Z^z with y = number of Y positions. Moving every X to the left
        costs (-1)^(z_j . x_l) for each pair j < l, so the product has sign
        r = sum r_j + sum_{j<l} z_j . x_l + (sum y_j - y) / 2 (mod 2), y = number of Y positions of the product.
        The pair term is computed per qubit column: prefix_parity(b) marks the rows below which an odd number of
        rows have Z at that qubit."""
        xacc = zacc = sumy = cross = 0
        for k in range(self.Q):
            a, b = self.sx[k] & J, self.sz[k] & J
            if a or b:
                xacc |= (popcount(a) & 1) << k
                zacc |= (popcount(b) & 1) << k
                sumy += popcount(a & b)
                if a and b:
                    cross ^= popcount(prefix_parity(b, self.Q) & a) & 1
        y = popcount(xacc & zacc)
        if (sumy - y) & 1:
            raise RuntimeError("non-Hermitian product of stabilizer rows")
        const = (popcount(self.r0 & J) + cross + ((sumy - y) >> 1)) & 1
        var = 0
        for j in bits(J):
            var ^= self.rv[j]
        return xacc, zacc, const | var

    def measure_z(self, q):
        """Measure Z_q; return the record as an affine expression; leave q decoupled in the state +-Z_q."""
        sx, sz, dx, dz, rv = self.sx, self.sz, self.dx, self.dz, self.rv
        anti = sx[q]
        if anti:  # random outcome
            p = (anti & -anti).bit_length() - 1
            bp = 1 << p
            px, pz = self.row(sx, sz, p)
            ms, md = anti ^ bp, dx[q]
            if ms:  # stabilizer rows i in ms <- stab_p * stab_i, exact sign
                # sign(stab_p stab_i) = r_p + r_i + z_p . x_i + (y_p + y_i - y_new) / 2. Per qubit k, with stab_p's bits
                # (a, b) fixed, y_i - y_new changes by +1 on rows with Y at k and by -1 on rows with Z (if a = 1, b = 0),
                # X (if a = 0, b = 1) or I (if a = b = 1); lo/hi hold that count mod 4, bit-sliced over the rows.
                t1 = lo = hi = 0
                for k in pz:
                    t1 ^= sx[k]
                spx, spz = set(px), set(pz)
                for k in spx | spz:
                    xk, zk = sx[k] & ms, sz[k] & ms
                    plus = xk & zk
                    if k in spx and k in spz:
                        minus = ms & ~(xk | zk)
                    elif k in spx:
                        minus = zk & ~xk
                    else:
                        minus = xk & ~zk
                    c = lo & plus
                    lo ^= plus
                    hi ^= c
                    c = ~lo & minus
                    lo ^= minus
                    hi ^= c
                ya = len(spx & spz) & 3
                if ya == 1 or ya == 3:
                    c = (lo & ms) if ya == 1 else (~lo & ms)
                    lo ^= ms
                    hi ^= c
                elif ya == 2:
                    hi ^= ms
                if lo & ms:
                    raise RuntimeError("anticommuting stabilizer rows")
                self.r0 ^= (ms if self.r0 & bp else 0) ^ ((t1 ^ hi) & ms)
                if rv[p]:
                    for i in bits(ms):
                        rv[i] ^= rv[p]
                for k in px:
                    sx[k] ^= ms
                for k in pz:
                    sz[k] ^= ms
            for k in px:
                dx[k] ^= md
            for k in pz:
                dz[k] ^= md
            for k in px:
                sx[k] &= ~bp
            for k in pz:
                sz[k] &= ~bp
            sz[q] |= bp
            self.r0 &= ~bp
            self.nvar += 1
            rv[p] = 1 << self.nvar
            rec = rv[p]
        else:  # deterministic: Z_q = +- product of the stabilizers whose destabilizer anticommutes with Z_q
            J = dx[q]
            xacc, zacc, rec = self.product(J)
            if J == 0 or xacc or zacc != 1 << q:
                raise RuntimeError("deterministic measurement is not in the stabilizer group")
            p = (J & -J).bit_length() - 1
            bp = 1 << p
            j2 = J ^ bp
            qx, qz = self.row(dx, dz, p)  # destab_i <- destab_i * destab_p for i in j2 (pairs stab_p <- product)
            for k in qx:
                dx[k] ^= j2
            for k in qz:
                dz[k] ^= j2
            px, pz = self.row(sx, sz, p)
            for k in px:
                sx[k] &= ~bp
            for k in pz:
                sz[k] &= ~bp
            sz[q] |= bp
            self.r0 = (self.r0 & ~bp) | (bp if rec & 1 else 0)
            rv[p] = rec & ~1
        # decouple: clear Z_q from every other row, destabilizer p := X_q. The state is now +-Z_q (row p) times the
        # rest; multiplying another row by row p keeps the group, destabilizer signs are unused, and once no other row
        # touches q, X_q is a valid destabilizer for row p. A later RZ / RX then just overwrites row p.
        others = sz[q] & ~bp
        if self.r0 & bp:
            self.r0 ^= others
        if rv[p]:
            for i in bits(others):
                rv[i] ^= rv[p]
        sz[q] = bp
        qx, qz = self.row(dx, dz, p)
        for k in qx:
            dx[k] &= ~bp
        for k in qz:
            dz[k] &= ~bp
        dx[q], dz[q] = bp, 0
        self.owner[q] = p
        return rec

    def prepare(self, q, basis):
        p = self.owner[q]
        bp = 1 << p
        if self.sx[q] or self.sz[q] != bp or self.dx[q] != bp or self.dz[q]:
            raise RuntimeError("preparing a qubit that is not decoupled")
        self.r0 &= ~bp
        self.rv[p] = 0
        if basis == "X":
            self.sx[q], self.sz[q], self.dx[q], self.dz[q] = bp, 0, 0, bp


def index_list(x, hi, what):
    want_list(x, what, max_len=hi)
    prev = -1
    for k in x:
        if not is_int(k) or not prev < k < hi:
            raise Reject(f"{what} must be a strictly increasing list of record indices below {hi}, got {short(x)}")
        prev = k
    return x


def parse_ops(ops, n, E, rmax):
    live, used, emitted = [False] * E, [False] * E, [False] * n
    out, nrec = [], 0

    def emitter(x, idx, need_live=True):
        e = want_int(x, 0, E - 1, f"ops[{idx}] emitter label")
        used[e] = True
        if need_live and not live[e]:
            raise Reject(f"ops[{idx}]: emitter {e} is not live (never prepared, or measured and not prepared again)")
        return e

    for idx, op in enumerate(ops):
        if not isinstance(op, list) or not op or not isinstance(op[0], str):
            raise Reject(f"ops[{idx}] must be a list that starts with an operation name")
        t = op[0]
        arity = 2 if t in ONE else 3 if t in ("CX", "CZ", "EMIT") else 4 if t in ("PX", "PZ") else None
        if arity is None:
            raise Reject(f"ops[{idx}]: unsupported operation {short(t)} (allowed: {', '.join(ONE)}, CX, CZ, PX, PZ, EMIT)")
        if len(op) != arity:
            raise Reject(f"ops[{idx}]: {t} must have exactly {arity - 1} operand(s)")
        if t in ("RZ", "RX"):
            e = emitter(op[1], idx, need_live=False)
            if live[e]:
                raise Reject(f"ops[{idx}]: {t} on live emitter {e}; a live emitter cannot be reset (measure it first)")
            live[e] = True
            out.append((t, e))
        elif t in ONE:
            e = emitter(op[1], idx)
            if t in ("MZ", "MX"):
                live[e] = False
                nrec += 1
                if nrec > rmax:
                    raise Reject(f"ops[{idx}]: more than measurements_max = {rmax} measurements")
            out.append((t, e))
        elif t in ("CX", "CZ"):
            a, b = emitter(op[1], idx), emitter(op[2], idx)
            if a == b:
                raise Reject(f"ops[{idx}]: {t} needs two distinct emitters")
            out.append((t, a, b))
        elif t in ("PX", "PZ"):
            e = emitter(op[1], idx)
            idxs = index_list(op[2], nrec, f"ops[{idx}] record list")
            out.append((t, e, tuple(idxs), want_int(op[3], 0, 1, f"ops[{idx}] value")))
        else:  # EMIT
            e = emitter(op[1], idx)
            v = want_int(op[2], 0, n - 1, f"ops[{idx}] photon label")
            if emitted[v]:
                raise Reject(f"ops[{idx}]: photon {v} is emitted a second time")
            emitted[v] = True
            out.append((t, e, v))
    missing = [v for v in range(n) if not emitted[v]]
    if missing:
        raise Reject(f"photon(s) never emitted: {short(missing)}")
    unused = [e for e in range(E) if not used[e]]
    if unused:
        raise Reject(f"declared emitter(s) never used in ops: {short(unused)}")
    return out, nrec


def parse_final(art, n, nrec):
    lcs = want_list(art["local_cliffords"], "local_cliffords", max_len=n, min_len=n)
    for v, c in enumerate(lcs):
        if not isinstance(c, dict):
            raise Reject(f"local_cliffords[{v}] must be an object")
        exact_keys(c, ["x", "z"], where=f"local_cliffords[{v}]")
        if c["x"] not in IMAGES or c["z"] not in IMAGES:
            raise Reject(f"local_cliffords[{v}]: images must be one of {', '.join(IMAGES)}")
        if c["x"][1] == c["z"][1]:
            raise Reject(f"local_cliffords[{v}]: the X and Z images must have different axes")
    cors = want_list(art["photon_corrections"], "photon_corrections", max_len=n, min_len=n)
    for v, c in enumerate(cors):
        if not isinstance(c, dict):
            raise Reject(f"photon_corrections[{v}] must be an object")
        exact_keys(c, ["x", "z", "x_const", "z_const"], where=f"photon_corrections[{v}]")
        index_list(c["x"], nrec, f"photon_corrections[{v}].x")
        index_list(c["z"], nrec, f"photon_corrections[{v}].z")
        want_int(c["x_const"], 0, 1, f"photon_corrections[{v}].x_const")
        want_int(c["z_const"], 0, 1, f"photon_corrections[{v}].z_const")
    return lcs, cors


def run(n, E, parsed):
    """Simulate the parsed ops. Returns (tableau, records, map variable -> record index, emitter-emitter gate count)."""
    T = Tableau(n + E)
    rec, var_rec = [], {}
    cee = 0
    for op in parsed:
        t = op[0]
        if t == "EMIT":
            T.cx(n + op[1], op[2])
            continue
        q = n + op[1]
        if t in ("CX", "CZ"):
            cee += 1
            (T.cx if t == "CX" else T.cz)(q, n + op[2])
        elif t == "H":
            T.h(q)
        elif t == "S":
            T.s(q)
        elif t == "S_DAG":
            T.s(q)
            T.z(q)
        elif t == "X":
            T.x(q)
        elif t == "Z":
            T.z(q)
        elif t in ("MZ", "MX"):
            if t == "MX":
                T.h(q)
            nv = T.nvar
            rec.append(T.measure_z(q))
            if T.nvar > nv:
                var_rec[T.nvar] = len(rec) - 1
        elif t in ("PX", "PZ"):
            a = op[3] ^ 1  # exponent = 1 iff XOR of the listed records equals value
            for k in op[2]:
                a ^= rec[k]
            T.flip(T.sz[q] if t == "PX" else T.sx[q], a)
        else:
            T.prepare(q, t[1])
    return T, rec, var_rec, cee


def verify(inst, art):
    return check(load_instance("emitter", inst), art)


def check(I, art):
    """verify() on an instance given as a loaded dict (tests pass small ad hoc graphs this way)."""
    n = I["n"]
    nbr = [0] * n
    for u, v in I["edges"]:
        nbr[u] |= 1 << v
        nbr[v] |= 1 << u
    exact_keys(art, ["emitters", "ops", "local_cliffords", "photon_corrections"])
    E = want_int(art["emitters"], 1, I["emitters_max"], "emitters")
    ops = want_list(art["ops"], "ops", max_len=I["ops_max"])
    parsed, nrec = parse_ops(ops, n, E, I["measurements_max"])
    lcs, cors = parse_final(art, n, nrec)

    T, rec, var_rec, cee = run(n, E, parsed)
    for v in range(n):
        T.clifford(v, lcs[v])
    for v in range(n):
        c = cors[v]
        a, b = c["x_const"], c["z_const"]
        for k in c["x"]:
            a ^= rec[k]
        for k in c["z"]:
            b ^= rec[k]
        T.flip(T.sz[v], a)
        T.flip(T.sx[v], b)

    for v in range(n):
        anti, J = T.sz[v], T.dz[v]
        for u in bits(nbr[v]):
            anti ^= T.sx[u]
            J ^= T.dx[u]
        if anti:
            raise Reject(f"final state is not the target: K_{v} = X_{v} Z(neighbours of {v}) is not a stabilizer "
                         "(wrong photon state, or photons still entangled with an emitter)")
        xacc, zacc, sign = T.product(J)
        if xacc != 1 << v or zacc != nbr[v]:
            raise RuntimeError("stabilizer decomposition mismatch")
        if sign & ~1:
            recs = [var_rec[j] for j in range(1, T.nvar + 1) if sign >> j & 1]
            raise Reject(f"final state is not the target for every outcome: the sign of K_{v} depends on the random "
                         f"measurement record(s) {short(recs)} (missing or wrong outcome correction)")
        if sign:
            raise Reject(f"final state is not the target: K_{v} has sign -1 for every outcome")

    return {"status": "accepted", "score": {"ee_gates": cee, "emitters": E},
            "detail": {"n": n, "emitters": E, "ee_gates": cee, "ops": len(ops), "measurements": nrec,
                       "random_measurements": T.nvar,
                       "check": f"all {n} graph stabilizers have sign +1 in the final state for every measurement outcome"}}