"""Board ft-stateprep (site key ftprep): fault-tolerant preparation of |0_L> of a small CSS code, ranked by an exact cost
(SPEC_v2_final section 4 with review_final_fable F1-F5).

Artifact {"qubits": Q, "ops": [...], "accept": [{"bits": [...], "value": 0|1}, ...]}. Wires 0..n-1 are the data outputs,
n..Q-1 ancillas. Ops: ["RZ",q] ["RX",q] (prepare |0> / |+>, only on an unused or measured wire), ["X",q] ["Z",q],
["CX",c,t] (c != t), ["MZ",q] ["MX",q] (destructive: one record bit, wire unavailable until prepared again),
["PX",q,bits,value] ["PZ",q,bits,value] (Pauli iff XOR of the earlier records in `bits` equals value). `accept` is a
conjunction of parity equations over the record. Record bit 0 means eigenvalue +1 (F1). All data wires live at the end.

Target: the state stabilized by +X(a), a in S_X = rowspan HX, and +Z(b), b in S_Z = rowspan(HZ, logical Z). Fault model:
every op is one location (RZ: X after; RX: Z after; X, Z, PX, PZ: X, Y or Z after the slot; CX: 15 two-qubit Paulis
after; MZ, MX: flipped record). Property: for every event of s <= t faults at distinct locations and every branch, an
accepted branch leaves the pure data state E|0_L> with mu_X(x) <= s and mu_Z(z) <= s (coset minimum weights of the X and
Z parts of E modulo S_X and S_Z separately), and the zero-plus-one-fault acceptance probability is >= 1/2.

Exact decision (no sampling, no floats):
1. Noiseless run: CSS stabilizer simulation in which every generator sign is an affine GF(2) function of the random
   measurement outcomes (one fresh variable per random measurement), so all branches are handled at once. Every
   acceptance equation must be identically satisfied, and every row of HX, HZ and logical Z must be a stabilizer of the
   output with sign identically +1. These n independent stabilizers act on the data only, so the data state is |0_L>
   and pure (unentangled from live ancillas) on every branch. An artifact failing this cannot meet the 1/2 floor (spec).
2. Faults: Pauli faults only shift constants. One pass propagates a Pauli frame whose entries are bitsets over basis
   fault variables (CX: X_c, Z_c, X_t, Z_t; X/Z/PX/PZ: X, Z; RZ: X; RX: Z; MZ/MX: record flip). A record flip that changes
   the parity of a conditional op's bits toggles that Pauli. Per alternative this gives, by linearity, its flip vector
   of the acceptance equations, and the commutation pattern of the data frame with the Z-type target rows (decides
   mu_X) and with the X-type rows (decides mu_Z). Signatures of an event are XORs; it is accepted iff its acceptance
   flips vanish (then on every branch, else on none).
3. mu_X(x) <= s iff the pattern of x lies in U_X(s) = {pattern of e : wt(e) <= s}, since S_X = (S_Z)^perp; same for Z.
   Singles are checked against U(1); for t = 2, pairs at distinct locations accepted together have equal acceptance
   flips (F5), so alternatives are grouped by flip vector and the distinct (X, Z) patterns of a group are XORed
   pairwise and tested against U(2), excluding pattern pairs that occur only at one common location.
Score: P = (1-p)^N + p (1-p)^(N-1) (A1 + A3/3 + A15/15), cost = C / P, both reduced hex fractions; then ancillas, then C.
"""
import os
import time
from collections import Counter
from fractions import Fraction
from itertools import combinations

import numpy as np

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

NALT = {"RZ": 1, "RX": 1, "MZ": 1, "MX": 1, "X": 3, "Z": 3, "PX": 3, "PZ": 3, "CX": 15}
NVAR = {"RZ": 1, "RX": 1, "MZ": 1, "MX": 1, "X": 2, "Z": 2, "PX": 2, "PZ": 2, "CX": 4}
ARITY = {"CX": 3, "PX": 4, "PZ": 4}
PAULI = "IXYZ"
HEX_MAX = 4096


class Budget(Exception):
    pass


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

def _bits(x, nrec, what):
    """Strictly increasing list of record indices that already exist (nrec records so far)."""
    want_list(x, what)  # length is bounded by the strictly-increasing check below
    prev = -1
    for b in x:
        if not is_int(b) or b < 0:
            raise Reject(f"{what} must contain record indices (nonnegative integers), got {short(b)}")
        if b >= nrec:
            raise Reject(f"{what} refers to record {b}, but only {nrec} records exist at that point (record lookahead)")
        if b <= prev:
            raise Reject(f"{what} must be strictly increasing")
        prev = b
    return tuple(x)


def parse(art, I):
    n = I["n"]
    exact_keys(art, ["qubits", "ops", "accept"])
    Q = want_int(art["qubits"], n, I["qubits_max"], "qubits")
    ops_l = want_list(art["ops"], "ops", max_len=I["ops_max"])
    acc_l = want_list(art["accept"], "accept", max_len=I["accept_equations_max"])
    state = [0] * Q   # 0 never prepared, 1 live, 2 measured
    nrec = 0
    ops = []

    def live(i, q, tag):
        if state[q] != 1:
            why = "was never prepared" if state[q] == 0 else "was measured and not prepared again"
            raise Reject(f"ops[{i}] {tag} uses wire {q}, which {why}")

    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 tag")
        tag = op[0]
        if tag not in NALT:
            raise Reject(f"ops[{i}]: unsupported operation {short(tag)}")
        if len(op) != ARITY.get(tag, 2):
            raise Reject(f"ops[{i}]: {tag} takes exactly {ARITY.get(tag, 2) - 1} arguments")
        q = want_int(op[1], 0, Q - 1, f"ops[{i}] wire")
        if tag in ("RZ", "RX"):
            if state[q] == 1:
                raise Reject(f"ops[{i}] {tag} {q}: reset of a live wire (prepare only an unused or measured wire)")
            state[q] = 1
            ops.append((tag, q))
        elif tag == "CX":
            t = want_int(op[2], 0, Q - 1, f"ops[{i}] target wire")
            if t == q:
                raise Reject(f"ops[{i}]: CX operands must be distinct")
            live(i, q, tag)
            live(i, t, tag)
            ops.append((tag, q, t))
        else:
            live(i, q, tag)
            if tag in ("MZ", "MX"):
                state[q] = 2
                nrec += 1
                if nrec > I["measurements_max"]:
                    raise Reject(f"more than {I['measurements_max']} measurements")
                ops.append((tag, q))
            elif tag in ("PX", "PZ"):
                bits = _bits(op[2], nrec, f"ops[{i}] condition bits")
                value = want_int(op[3], 0, 1, f"ops[{i}] condition value")
                ops.append((tag, q, bits, value))
            else:
                ops.append((tag, q))
    for d in range(n):
        if state[d] != 1:
            raise Reject(f"data wire {d} {'was never prepared' if state[d] == 0 else 'is measured'} at the end; "
                         f"all data outputs 0..{n - 1} must be live")
    accept = []
    for j, e in enumerate(acc_l):
        if not isinstance(e, dict):
            raise Reject(f"accept[{j}] must be an object")
        exact_keys(e, ["bits", "value"], where=f"accept[{j}]")
        accept.append((_bits(e["bits"], nrec, f"accept[{j}].bits"), want_int(e["value"], 0, 1, f"accept[{j}].value")))
    return Q, ops, accept, nrec


def _xor(vals, bits):
    a = 0
    for b in bits:
        a ^= vals[b]
    return a


# ---------------------------------------------------------------- 1. noiseless run with affine signs

def noiseless(ops, accept, nrec, XT, ZT, names_x, names_z):
    """Reject unless every acceptance equation holds and every target row is a +1 stabilizer on every noiseless branch.
    Generators are ints (support << SB) | sign, sign = affine function: bit 0 constant, bit k >= 1 fresh variable k."""
    SB = nrec + 1
    MASK = (1 << SB) - 1
    G = {"X": [], "Z": []}
    rec = []
    fresh = 0
    for op in ops:
        tag = op[0]
        if tag in ("RZ", "RX"):
            G[tag[1]].append(1 << (op[1] + SB))
        elif tag == "CX":
            cb, tb = 1 << (op[1] + SB), 1 << (op[2] + SB)
            G["X"] = [g ^ tb if g & cb else g for g in G["X"]]     # X_c -> X_c X_t
            G["Z"] = [g ^ cb if g & tb else g for g in G["Z"]]     # Z_t -> Z_c Z_t
        elif tag in ("X", "Z", "PX", "PZ"):
            # applied indicator (affine): 1 for X/Z; 1 + value + XOR(records) for PX/PZ
            s = 1 if len(op) == 2 else 1 ^ op[3] ^ _xor(rec, op[2])
            if s:
                qb = 1 << (op[1] + SB)
                kind = "Z" if tag[-1] == "X" else "X"              # an X flips the Z-type generators on q
                G[kind] = [g ^ s if g & qb else g for g in G[kind]]
        else:
            qb = 1 << (op[1] + SB)
            same, other = ("Z", "X") if tag == "MZ" else ("X", "Z")
            anti = [k for k, g in enumerate(G[other]) if g & qb]
            if anti:                                               # random outcome: fresh variable
                g0 = G[other][anti[0]]
                for k in anti[1:]:
                    G[other][k] ^= g0
                del G[other][anti[0]]
                fresh += 1
                out = 1 << fresh
                G[same] = [g ^ qb ^ out if g & qb else g for g in G[same]]   # times (-1)^out Z_q, then drop q
            else:                                                  # deterministic: Z_q is in the Z-type span
                basis = rref(G[same])
                p = op[1] + SB
                row = basis.get(p)
                if row is None or row >> SB != 1 << op[1]:
                    raise RuntimeError("internal: deterministic measurement has no stabilizer")
                out = row & MASK
                G[same] = [g for k, g in basis.items() if k != p]
            rec.append(out)
    for j, (bits, value) in enumerate(accept):
        a = value ^ _xor(rec, bits)
        if a == 1:
            raise Reject(f"accept[{j}] fails on every noiseless branch: the noiseless circuit is never accepted")
        if a:
            raise Reject(f"accept[{j}] depends on random measurement outcomes, so it fails on half of the noiseless "
                         "branches; the acceptance probability floor 1/2 cannot then be met")
    for kind, rows, names in (("X", XT, names_x), ("Z", ZT, names_z)):
        basis = rref(G[kind])
        for h, name in zip(rows, names):
            v = reduce(h << SB, basis)
            if v >> SB:
                raise Reject(f"noiseless output is not |0_L>: {name} ({kind} type) is not a stabilizer of the output "
                             "(data outside the target state or entangled with live ancillas)")
            if v == 1:
                raise Reject(f"noiseless output is not |0_L>: {name} ({kind} type) has eigenvalue -1")
            if v:
                raise Reject(f"noiseless output is not |0_L>: the sign of {name} ({kind} type) depends on random "
                             "measurement outcomes (some branch is not corrected)")
    return fresh


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

def _transpose(rows, V):
    """rows: R ints over V bits -> V ints over R bits (bit r of out[v] = bit v of rows[r])."""
    if not rows:
        return [0] * V
    nb = (V + 7) // 8
    M = np.frombuffer(b"".join(r.to_bytes(nb, "little") for r in rows), dtype=np.uint8).reshape(len(rows), nb)
    B = np.unpackbits(M, axis=1, bitorder="little")[:, :V]
    P = np.packbits(np.ascontiguousarray(B.T), axis=1, bitorder="little")
    return [int.from_bytes(P[v].tobytes(), "little") for v in range(V)]


def signatures(Q, ops, accept, n, XT, ZT):
    """Per op: list of (alternative name, acceptance flips, X pattern, Z pattern)."""
    fx, fz = [0] * Q, [0] * Q
    flip = []
    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)
        elif tag in ("MZ", "MX"):
            flip.append((fx[q] if tag == "MZ" else fz[q]) ^ (1 << v))
            fx[q] = fz[q] = 0
        else:
            if tag == "PX":
                fx[q] ^= _xor(flip, op[2])
            elif tag == "PZ":
                fz[q] ^= _xor(flip, op[2])
            fx[q] ^= 1 << v
            fz[q] ^= 1 << (v + 1)
        v += NVAR[tag]
    V = v
    acc_rows = [_xor(flip, bits) for bits, _ in accept]
    xp_rows = [_xor(fx, [d for d in range(n) if h >> d & 1]) for h in ZT]   # X part vs Z-type target rows
    zp_rows = [_xor(fz, [d for d in range(n) if h >> d & 1]) for h in XT]   # Z part vs X-type target rows
    acc_v, xp_v, zp_v = _transpose(acc_rows, V), _transpose(xp_rows, V), _transpose(zp_rows, V)
    out = []
    for i, op in enumerate(ops):
        tag, v = op[0], first[i]
        if NALT[tag] == 1:
            comb = [("flip" if tag[0] == "M" else ("X" if tag == "RZ" else "Z"), (v,))]
        elif NALT[tag] == 3:
            comb = [("X", (v,)), ("Y", (v, v + 1)), ("Z", (v + 1,))]
        else:
            comb = []
            for a in range(1, 16):
                ic, it = a >> 2, a & 3
                vs = tuple(u for u, on in ((v, ic in (1, 2)), (v + 1, ic in (2, 3)), (v + 2, it in (1, 2)),
                                           (v + 3, it in (2, 3))) if on)
                comb.append((PAULI[ic] + PAULI[it], vs))
        alts = []
        for name, vs in comb:
            a = x = z = 0
            for u in vs:
                a ^= acc_v[u]
                x ^= xp_v[u]
                z ^= zp_v[u]
            alts.append((name, a, x, z))
        out.append(alts)
    return out


# ---------------------------------------------------------------- 3. syndrome sets and the event checks

def _cols(rows, n):
    col = [0] * n
    for r, h in enumerate(rows):
        for d in range(n):
            if h >> d & 1:
                col[d] |= 1 << r
    return col


def _uset(col, s):
    U = {0}
    if s >= 1:
        U |= set(col)
    if s >= 2:
        U |= {a ^ b for a, b in combinations(col, 2)}
    return U


class Member:
    """Vectorized membership in a set of patterns of the given bit width."""

    def __init__(self, U, width):
        self.table = None
        if width <= 22:
            self.table = np.zeros(1 << width, dtype=bool)
            self.table[np.array(sorted(U), dtype=np.int64)] = True
        else:
            self.sorted = np.array(sorted(U), dtype=np.uint64)

    def __call__(self, arr):
        if self.table is not None:
            return self.table[arr.astype(np.int64)]
        idx = np.minimum(np.searchsorted(self.sorted, arr), len(self.sorted) - 1)
        return self.sorted[idx] == arr


def mu_text(pattern, col, s, wmax=4):
    """Exact coset minimum weight of an error with this pattern if <= wmax (by enumeration), as text vs the bound s."""
    n = len(col)
    for w in range(0, min(wmax, n) + 1):
        for c in combinations(range(n), w):
            a = 0
            for d in c:
                a ^= col[d]
            if a == pattern:
                return f"{w}" + (f" > {s}" if w > s else f" <= {s}")
    return f"> {wmax} (> {s})"


def describe(ops, i):
    return short([list(x) if isinstance(x, tuple) else x for x in ops[i]], 50)


def verify(inst, art, time_budget=None):
    I = load_instance("ftprep", inst)
    budget = time_budget if time_budget is not None else float(os.environ.get("QARENA_CPU", I["cpu_seconds"])) - 60
    deadline = time.process_time() + budget          # the board's budget is CPU seconds
    t0 = time.time()
    n, t = I["n"], I["t"]
    bits_of = lambda sup: sum(1 << q for q in sup)
    XT = [bits_of(r) for r in I["HX"]]
    ZT = [bits_of(r) for r in I["HZ"]] + [bits_of(I["logical_z"])]
    names_x = [f"HX[{r}]" for r in range(len(I["HX"]))]
    names_z = [f"HZ[{r}]" for r in range(len(I["HZ"]))] + ["logical Z"]
    assert len(ZT) <= 64 and len(XT) <= 64

    Q, ops, accept, nrec = parse(art, I)
    N = len(ops)
    f = [NALT[op[0]] for op in ops]
    F = sum(f)
    K = 1 + F + (t == 2) * (F * F - sum(x * x for x in f)) // 2
    if K > I["fault_cases_max"]:
        raise Reject(f"{K} fault cases exceed the instance limit {I['fault_cases_max']}")

    fresh = noiseless(ops, accept, nrec, XT, ZT, names_x, names_z)
    try:
        if time.process_time() > deadline:
            raise Budget
        sig = signatures(Q, ops, accept, n, XT, ZT)
        colx, colz = _cols(ZT, n), _cols(XT, n)

        def fail(s, events, x, z):
            ev = "; ".join(f"op {i} {describe(ops, i)} alternative {name}" for i, name in events)
            raise Reject(f"fault tolerance violated by an accepted {s}-fault event: {ev}. Data error has "
                         f"mu_X = {mu_text(x, colx, s)}, mu_Z = {mu_text(z, colz, s)} (both must be <= {s})")

        # single faults (s = 1) and the A counts
        UX1, UZ1 = _uset(colx, 1), _uset(colz, 1)
        A = {1: 0, 3: 0, 15: 0}
        for i, alts in enumerate(sig):
            for name, a, x, z in alts:
                if a == 0:
                    if x not in UX1 or z not in UZ1:
                        fail(1, [(i, name)], x, z)
                    A[len(alts)] += 1
        pair_events = 0
        if t == 2:
            UX2, UZ2 = _uset(colx, 2), _uset(colz, 2)
            inX, inZ = Member(UX2, len(ZT)), Member(UZ2, len(XT))
            groups = {}
            for i, alts in enumerate(sig):
                for k, (name, a, x, z) in enumerate(alts):
                    groups.setdefault(a, []).append((i, k))
            for a, mem in groups.items():
                cnt = Counter(i for i, _ in mem)
                tot = len(mem)
                pair_events += (tot * tot - sum(c * c for c in cnt.values())) // 2
                if len(cnt) < 2:
                    continue
                # distinct (X, Z) patterns: single location or -1, representative, representative elsewhere
                d = {}
                for i, k in mem:
                    _, _, x, z = sig[i][k]
                    e = d.get((x, z))
                    if e is None:
                        d[(x, z)] = [i, (i, k), None]
                    elif e[0] != -1 and i != e[0]:
                        e[0], e[2] = -1, (i, k)
                keys = list(d)
                m = len(keys)
                if m < 2:
                    continue
                xs = np.array([x for x, _ in keys], dtype=np.uint64)
                zs = np.array([z for _, z in keys], dtype=np.uint64)
                single = np.array([d[key][0] for key in keys], dtype=np.int64)
                for u in range(m - 1):
                    if (u & 255) == 0 and time.process_time() > deadline:
                        raise Budget
                    bad = ~(inX(xs[u + 1:] ^ xs[u]) & inZ(zs[u + 1:] ^ zs[u]))
                    if single[u] >= 0:
                        bad &= single[u + 1:] != single[u]
                    if bad.any():
                        w = u + 1 + int(np.argmax(bad))
                        eu, ew = d[keys[u]], d[keys[w]]
                        ru, rw = eu[1], ew[1]
                        if ru[0] == rw[0]:
                            if eu[2] is not None:
                                ru = eu[2]
                            else:
                                rw = ew[2]
                        evs = sorted([(ru[0], sig[ru[0]][ru[1]][0]), (rw[0], sig[rw[0]][rw[1]][0])])
                        fail(2, evs, keys[u][0] ^ keys[w][0], keys[u][1] ^ keys[w][1])
    except Budget:
        return {"status": "error", "detail": {"reason": "fault enumeration did not finish within the time limit "
                                                        "(not ranked)", "fault_cases": K}}

    C = sum(1 for op in ops if op[0] == "CX")
    p = Fraction(I["p_num"], I["p_den"])
    P = (1 - p) ** N + p * (1 - p) ** (N - 1) * (A[1] + Fraction(A[3], 3) + Fraction(A[15], 15))
    if 2 * P.numerator < P.denominator:
        raise Reject(f"acceptance probability lower bound {float(P):.6f} is below 1/2 (exact comparison of the "
                     f"fractions); the circuit has N = {N} operations and A1 + A3/3 + A15/15 = "
                     f"{float(A[1] + Fraction(A[3], 3) + Fraction(A[15], 15)):.4f} tolerated single faults")
    cost = Fraction(C) / P
    hx = lambda fr: {"num": format(fr.numerator, "x"), "den": format(fr.denominator, "x")}
    sc_cost, sc_p = hx(cost), hx(P)
    if max(len(s) for s in list(sc_cost.values()) + list(sc_p.values())) > HEX_MAX:
        return {"status": "error", "detail": {"reason": "score fraction longer than 4096 hex digits"}}
    return {"status": "accepted",
            "score": {"cost": sc_cost, "ancillas": Q - n, "cx": C, "p_accept_lower": sc_p, "locations": N,
                      "a1": A[1], "a3": A[3], "a15": A[15]},
            "detail": {"n": n, "t": t, "qubits": Q, "locations": N, "cx": C, "measurements": nrec,
                       "random_measurements": fresh, "accept_equations": len(accept),
                       "alternatives": F, "fault_cases": K,
                       "accepted_single_alternatives": A[1] + A[3] + A[15],
                       "rejected_single_alternatives": F - A[1] - A[3] - A[15],
                       "accepted_pair_events": pair_events if t == 2 else None,
                       "cost_decimal": f"{float(cost):.12g}", "p_accept_lower_decimal": f"{float(P):.12g}",
                       "note": "every noiseless branch is accepted with output |0_L>; every event of at most t faults "
                               "was checked exactly (decimals are for display only)",
                       "verify_seconds": round(time.time() - t0, 2)}}