"""Board flag-check (site key flagcheck): exact parity instruments with flags and two-fault hook detection
(SPEC_v3_final s4, which folds in review_v3_fable F1-F3 and G1-G4).

Instance: w data wires, P = Z^w, fault order t (2 on every board instance), caps on declared ancillas, ops, CX and
measurements. Artifact {"ancillas": A, "ops": [...], "syndrome": j}. Data wires 0..w-1 are live throughout and are never
prepared or measured; ancilla wires w..w+A-1 start inactive and must be inactive at the end. Ops: ["RZ",q] ["RX",q] prepare an
inactive ancilla in |0> / |+>; ["CX",c,t] acts on two distinct live wires; ["MZ",q] ["MX",q] measure a live ancilla, append one
record bit (0 = eigenvalue +1) and make it inactive. Record `syndrome` is the parity outcome; every other record is a flag.

Exact decision (no sampling, no floats):
1. Ideal instrument. One symbolic run on the Choi state (Bell pairs D_i R_i; the references R_i are never touched). The state
   stays CSS, so it is kept as X-type and Z-type generator lists; a generator is an int (support << SB) | sign, where the sign
   is an affine GF(2) form (bit 0 constant, bit k >= 1 the outcome of the k-th random measurement). One run covers every
   branch. Required: every flag record is identically 0; the syndrome record has a nonzero random part (unbiased); the final
   state contains +Z_Di Z_Ri (all i), +X_D0 X_R0 X_Di X_Ri (i >= 1) and (-1)^syndrome prod_i Z_Di. These 2w independent
   signed generators fix the conditional Choi state, so every branch applies Pi_s up to a phase: the instrument is exactly
   rho -> Pi_s rho Pi_s with Pi_s = (I + (-1)^s P) / 2.
2. Faults (frozen location set: every listed op is one location; no idle locations). CX: the 15 nonidentity two-qubit Paulis
   after the gate; RZ: X after; RX: Z after; MZ / MX: the record bit flipped. One forward Pauli-frame pass over basis
   variables (CX: X_c, Z_c, X_t, Z_t; RZ: X; RX: Z; MZ / MX: flip) gives, for every variable, its flag-flip vector (the
   syndrome record excluded) and its final data masks (x, z); measurement discards the wire's frame and preparation starts a
   cleared one. By linearity an alternative's signature is the XOR of its variables' and an event's the XOR of its
   alternatives'. An event of v = 1 or 2 faults that flips no flag must leave min(popcount(x|z), popcount(x|(z^D))) <= v,
   D = all data wires (wt(E) and wt(E P)). Pairs are alternatives at distinct locations. Only equal flag vectors cancel, so
   alternatives are grouped by flag vector; inside a group the distinct (x, z) patterns are compared pairwise, skipping only a
   pattern pair whose occurrences all sit at one common location.
Score (all minimized, then submission number): cx, ancillas (declared), measurements.
Cost: F = 15 C + P + M <= 1984 alternatives, fewer than 2 million pairs; well under a second at the caps.
"""
import os
import time

from common import Reject, exact_keys, want_int, want_list, is_int, load_instance, rref, reduce, short

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

ARITY = {"RZ": 2, "RX": 2, "CX": 3, "MZ": 2, "MX": 2}
RZ, RX, CX, MZ, MX = range(5)
CODE = {"RZ": RZ, "RX": RX, "CX": CX, "MZ": MZ, "MX": MX}
NVAR = (1, 1, 4, 1, 1)
NALT = (1, 1, 15, 1, 1)
SINGLE_NAME = ("X", "Z", None, "flip", "flip")
PAULI = "IXYZ"


class Budget(Exception):
    pass


def check_instance(I):
    """The instance must be self-consistent; anything else is an internal inconsistency (status error, never a verdict)."""
    ok = (isinstance(I, dict)
          and all(is_int(I.get(k)) for k in ("data_qubits", "fault_order", "ancillas_max", "ops_max", "cx_max",
                                              "measurements_max", "cpu_seconds"))
          and 1 <= I["data_qubits"] <= 32 and I["fault_order"] in (1, 2) and I["pauli"] == "Z" * I["data_qubits"]
          and I["ancillas_max"] >= 0 and I["ops_max"] >= 0 and I["cx_max"] >= 0 and I["measurements_max"] >= 0
          and I["cpu_seconds"] > 0 and I.get("connectivity") == "all-to-all" and I.get("noise") == "operation-pauli")
    if not ok:
        raise RuntimeError("instance file is inconsistent")


# ---------------------------------------------------------------- parsing and wire lifecycle

def parse(art, I):
    """Strict parse plus the lifecycle. Returns (A, ops, nrec, syndrome, mops) with ops as tuples (code, q[, t]) and mops[j]
    the op index of record j."""
    w = I["data_qubits"]
    if not isinstance(art, dict):
        raise Reject("artifact must be a JSON object")
    exact_keys(art, ["ancillas", "ops", "syndrome"])
    A = want_int(art["ancillas"], 0, I["ancillas_max"], "ancillas")
    ops_l = want_list(art["ops"], "ops", max_len=I["ops_max"])
    W = w + A
    live = [True] * w + [False] * A
    ops, mops = [], []
    ncx = 0
    for i, op in enumerate(ops_l):
        if not isinstance(op, list) or not op or not isinstance(op[0], str):
            raise Reject(f"ops[{i}] must be a list that starts with an operation name")
        tag = op[0]
        if tag not in ARITY:
            raise Reject(f"ops[{i}]: unsupported operation {short(tag)} (allowed: RZ, RX, CX, MZ, MX)")
        if len(op) != ARITY[tag]:
            raise Reject(f"ops[{i}]: {tag} takes exactly {ARITY[tag] - 1} wire argument(s)")
        q = want_int(op[1], 0, W - 1, f"ops[{i}] wire")
        if tag == "CX":
            t = want_int(op[2], 0, W - 1, f"ops[{i}] target wire")
            if t == q:
                raise Reject(f"ops[{i}]: CX control and target must be different wires")
            for x in (q, t):
                if not live[x]:
                    raise Reject(f"ops[{i}]: CX uses ancilla {x}, which is not live (never prepared, or measured "
                                 "and not prepared again)")
            ncx += 1
            if ncx > I["cx_max"]:
                raise Reject(f"more than {I['cx_max']} CX operations")
            ops.append((CX, q, t))
            continue
        if q < w:
            raise Reject(f"ops[{i}]: {tag} on data wire {q}; data wires 0..{w - 1} can never be prepared or measured")
        if tag in ("RZ", "RX"):
            if live[q]:
                raise Reject(f"ops[{i}]: {tag} {q} prepares a live ancilla (prepare only an inactive ancilla)")
            live[q] = True
        else:
            if not live[q]:
                raise Reject(f"ops[{i}]: {tag} {q} measures an inactive ancilla")
            live[q] = False
            mops.append(i)
            if len(mops) > I["measurements_max"]:
                raise Reject(f"more than {I['measurements_max']} measurements")
        ops.append((CODE[tag], q))
    for q in range(w, W):
        if live[q]:
            raise Reject(f"ancilla {q} is still live at the end; every ancilla must be measured before the gadget ends")
    nrec = len(mops)
    if nrec == 0:
        raise Reject("the gadget has no measurement, so it has no syndrome record")
    syn = want_int(art["syndrome"], 0, nrec - 1, "syndrome (an existing record index)")
    return A, ops, nrec, syn, mops


# ---------------------------------------------------------------- 1. the ideal instrument

def ideal(w, W, ops, nrec, syn, mops):
    """Symbolic Choi-state run (see the module docstring). Raises Reject unless the noiseless instrument is exactly
    rho -> Pi_s rho Pi_s with all flags 0. Support positions: data and ancilla wire q at q, reference R_i at W + i."""
    SB = nrec + 1
    SMASK = (1 << SB) - 1
    Xg = [((1 << i) | (1 << (W + i))) << SB for i in range(w)]
    Zg = list(Xg)
    rec = []
    fresh = 0
    for op in ops:
        tag, q = op[0], op[1]
        if tag == RZ:
            Zg.append(1 << (q + SB))
        elif tag == RX:
            Xg.append(1 << (q + SB))
        elif tag == CX:
            cb, tb = 1 << (q + SB), 1 << (op[2] + SB)
            Xg = [g ^ tb if g & cb else g for g in Xg]     # X_c -> X_c X_t
            Zg = [g ^ cb if g & tb else g for g in Zg]     # Z_t -> Z_c Z_t
        else:
            qb = 1 << (q + SB)
            same, other = (Zg, Xg) if tag == MZ else (Xg, Zg)
            anti = [k for k, g in enumerate(other) if g & qb]
            if anti:                                       # random outcome: a fresh variable
                fresh += 1
                out = 1 << fresh
                g0 = other[anti[0]]
                for k in anti[1:]:
                    other[k] ^= g0
                del other[anti[0]]
                for k, g in enumerate(same):               # times (-1)^out P_q, which removes the measured wire
                    if g & qb:
                        same[k] = g ^ qb ^ out
            else:                                          # deterministic: P_q lies in the same-type span
                basis = rref(same)
                row = basis.get(q + SB)
                if row is None or row >> SB != 1 << q:
                    raise RuntimeError("internal: deterministic measurement without a stabilizer")
                out = row & SMASK
                same[:] = [g for p, g in basis.items() if p != q + SB]
            rec.append(out)
    anc = ((1 << W) - (1 << w)) << SB
    if len(Xg) + len(Zg) != 2 * w or any(g & anc for g in Xg + Zg):
        raise RuntimeError("internal: the final state is not a 2w-qubit data/reference state")

    for j, r in enumerate(rec):
        if j != syn and r:
            how = "always 1" if r == 1 else "1 on half of the noiseless runs"
            raise Reject(f"ideal instrument: flag record {j} (ops[{mops[j]}]) must be 0 on every noiseless run, but it "
                         f"is {how}")
    s = rec[syn]
    if not s >> 1:
        raise Reject(f"ideal instrument: the syndrome record {syn} (ops[{mops[syn]}]) is deterministic (always {s}), so "
                     "the gadget does not measure Z^w")
    bx, bz = rref(Xg), rref(Zg)
    d = lambda i: 1 << (i + SB)          # data qubit D_i
    ref = lambda i: 1 << (W + i + SB)    # reference qubit R_i

    def need(basis, g, name, meaning):
        v = reduce(g, basis)
        if v >> SB:
            raise Reject(f"ideal instrument: {name} is not a stabilizer of the final data/reference (Choi) state: {meaning}")
        if v:
            raise Reject(f"ideal instrument: {name} has the wrong sign on some noiseless run: {meaning}")

    for i in range(w):
        need(bz, d(i) | ref(i), f"+Z(D{i})Z(R{i})", f"the gadget does not leave the Z value of data qubit {i} unchanged")
    for i in range(1, w):
        need(bx, d(0) | ref(0) | d(i) | ref(i), f"+X(D0)X(R0)X(D{i})X(R{i})",
             f"the gadget disturbs more than the parity of data qubits 0 and {i}")
    need(bz, sum(d(i) for i in range(w)) | s, "(-1)^syndrome Z(D0)...Z(D%d)" % (w - 1),
         "the syndrome record is not the measured parity")
    return fresh


# ---------------------------------------------------------------- 2. fault signatures

def _spread(vec, bit, out):
    """For every set bit u of vec: out[u] |= bit."""
    while vec:
        low = vec & -vec
        out[low.bit_length() - 1] |= bit
        vec ^= low


def alternatives(w, W, ops, syn):
    """Every fault alternative as (op index, name, flag flips, data x mask, data z mask)."""
    fx, fz = [0] * W, [0] * W
    flips, first = [], []
    v = 0
    for op in ops:
        tag, q = op[0], op[1]
        first.append(v)
        if tag == RZ:
            fx[q], fz[q] = 1 << v, 0
        elif tag == RX:
            fx[q], fz[q] = 0, 1 << v
        elif tag == CX:
            t = op[2]
            fx[t] ^= fx[q]
            fz[q] ^= fz[t]
            fx[q] ^= 1 << v
            fz[q] ^= 1 << (v + 1)
            fx[t] ^= 1 << (v + 2)
            fz[t] ^= 1 << (v + 3)
        else:
            flips.append((fx[q] if tag == MZ else fz[q]) ^ (1 << v))
            fx[q] = fz[q] = 0
        v += NVAR[tag]
    V = v
    fl, xs, zs = [0] * V, [0] * V, [0] * V
    k = 0
    for j, f in enumerate(flips):
        if j != syn:
            _spread(f, 1 << k, fl)
            k += 1
    for dq in range(w):
        _spread(fx[dq], 1 << dq, xs)
        _spread(fz[dq], 1 << dq, zs)
    out = []
    for i, op in enumerate(ops):
        v = first[i]
        if op[0] != CX:
            out.append((i, SINGLE_NAME[op[0]], fl[v], xs[v], zs[v]))
            continue
        for a in range(1, 16):
            pc, pt = a >> 2, a & 3
            f = x = z = 0
            for u, on in ((v, pc in (1, 2)), (v + 1, pc in (2, 3)), (v + 2, pt in (1, 2)), (v + 3, pt in (2, 3))):
                if on:
                    f ^= fl[u]
                    x ^= xs[u]
                    z ^= zs[u]
            out.append((i, PAULI[pc] + PAULI[pt], f, x, z))
    return out


def weight(x, z, D):
    """min(wt(E), wt(E Z^w)) for E = X(x) Z(z) on the data."""
    return min(popcount(x | z), popcount(x | (z ^ D)))


def pauli_text(x, z, w):
    return "".join("IXZY"[(x >> d & 1) | (z >> d & 1) << 1] for d in range(w))


# ---------------------------------------------------------------- 3. the event checks

def check_events(w, ops_l, alts, t, deadline_check):
    """Raise Reject at the first violating event (singles, then pairs). Returns (unflagged singles, pair events, unflagged
    pair events)."""
    D = (1 << w) - 1

    def fail(v, evs, x, z):
        ev = " and ".join(f"ops[{i}] {short(ops_l[i], 40)} fault {name}" for i, name in evs)
        what = "single fault" if v == 1 else "two-fault event"
        raise Reject(f"fault condition violated by the {what} {ev}: no flag fires and the data error is "
                     f"E = {pauli_text(x, z, w)} (min(wt(E), wt(E Z^w)) = {weight(x, z, D)} > {v})")

    un1 = 0
    for i, name, f, x, z in alts:
        if f == 0:
            un1 += 1
            if weight(x, z, D) > 1:
                fail(1, [(i, name)], x, z)
    groups = {}
    for k, a in enumerate(alts):
        groups.setdefault(a[2], []).append(k)
    F = len(alts)
    per_loc = {}
    for a in alts:
        per_loc[a[0]] = per_loc.get(a[0], 0) + 1
    pairs = (F * F - sum(c * c for c in per_loc.values())) // 2
    un2 = 0
    for f, mem in groups.items():
        cnt = {}
        for k in mem:
            cnt[alts[k][0]] = cnt.get(alts[k][0], 0) + 1
        un2 += (len(mem) ** 2 - sum(c * c for c in cnt.values())) // 2
        if t < 2 or len(cnt) < 2:
            continue
        # distinct (x, z) patterns: [only location or -1 when several, representative, representative at another location]
        pat = {}
        for k in mem:
            i, _, _, x, z = alts[k]
            e = pat.get((x, z))
            if e is None:
                pat[(x, z)] = [i, k, None]
            elif e[0] != -1 and i != e[0]:
                e[0], e[2] = -1, k
        keys = list(pat)
        locs = [pat[key][0] for key in keys]
        m = len(keys)
        for u in range(m - 1):
            if not u & 63:
                deadline_check()
            xu, zu = keys[u]
            lu = locs[u]
            for v in range(u + 1, m):
                xv, zv = keys[v]
                x, z = xu ^ xv, zu ^ zv
                if popcount(x | z) > 2 and popcount(x | (z ^ D)) > 2 and (lu == -1 or locs[v] != lu):
                    eu, ev = pat[keys[u]], pat[keys[v]]
                    ku, kv = eu[1], ev[1]
                    if alts[ku][0] == alts[kv][0]:                 # pick representatives at distinct locations
                        if eu[2] is not None:
                            ku = eu[2]
                        else:
                            kv = ev[2]
                    a, b = sorted((ku, kv))
                    fail(2, [(alts[a][0], alts[a][1]), (alts[b][0], alts[b][1])], x, z)
    return un1, pairs, un2


# ---------------------------------------------------------------- entry points

def verify(inst, art, time_budget=None):
    return check(load_instance("flagcheck", inst), art, time_budget)


def check(I, art, time_budget=None):
    """verify() on an instance given as a loaded dict (tests pass small ad hoc instances this way)."""
    check_instance(I)
    w, t = I["data_qubits"], I["fault_order"]
    budget = time_budget if time_budget is not None else max(1.0, float(os.environ.get("QARENA_CPU", I["cpu_seconds"])) - 30)
    t0 = time.process_time()

    def deadline_check():
        if time.process_time() - t0 > budget:
            raise Budget

    A, ops, nrec, syn, mops = parse(art, I)
    W = w + A
    ideal(w, W, ops, nrec, syn, mops)
    try:
        deadline_check()
        alts = alternatives(w, W, ops, syn)
        un1, pairs, un2 = check_events(w, art["ops"], alts, t, deadline_check)
    except Budget:
        return {"status": "error", "detail": {"reason": "fault enumeration did not finish within the CPU budget (not ranked)"}}
    C = sum(1 for op in ops if op[0] == CX)
    return {"status": "accepted",
            "score": {"cx": C, "ancillas": A, "measurements": nrec},
            "detail": {"data_qubits": w, "fault_order": t, "cx": C, "ancillas": A, "measurements": nrec, "flags": nrec - 1,
                       "locations": len(ops), "alternatives": len(alts), "pair_events": pairs if t == 2 else None,
                       "unflagged_single_alternatives": un1, "unflagged_pair_events": un2 if t == 2 else None,
                       "check": "noiseless: every flag 0, unbiased syndrome, exact signed Choi state of Pi_s; every "
                                + ("single fault and every pair of faults at distinct operations" if t == 2 else
                                   "single fault") + " with no flag fired leaves min(wt(E), wt(E Z^w)) <= number of faults"}}


def analyze(I, art):
    """Test helper (not used by verify): parse and ideal check as in check(), then exact counts over the whole fault
    model by brute force, without early exit: locations, alternatives, distinct-location pairs, unflagged and violating
    singles and pairs."""
    check_instance(I)
    w = I["data_qubits"]
    A, ops, nrec, syn, mops = parse(art, I)
    ideal(w, w + A, ops, nrec, syn, mops)
    alts = alternatives(w, w + A, ops, syn)
    D = (1 << w) - 1
    un1 = sum(1 for a in alts if a[2] == 0)
    bad1 = sum(1 for a in alts if a[2] == 0 and weight(a[3], a[4], D) > 1)
    pairs = un2 = bad2 = 0
    for p in range(len(alts)):
        i, _, f, x, z = alts[p]
        for q in range(p + 1, len(alts)):
            j, _, g, y, u = alts[q]
            if i == j:
                continue
            pairs += 1
            if f == g:
                un2 += 1
                if weight(x ^ y, z ^ u, D) > 2:
                    bad2 += 1
    return {"locations": len(ops), "alternatives": len(alts), "pairs": pairs, "unflagged1": un1, "bad1": bad1,
            "unflagged2": un2, "bad2": bad2}