"""Board pauli-network (site key pnet): compile a fixed ordered list of symbolic Pauli rotations into a circuit of Clifford
gates and single-qubit Z rotations with few CX gates (SPEC_v6 s3 with review_v6_fable N1-N2).

Target U = R_{P_{M-1}}(t_{M-1}) ... R_{P_0}(t_0), R_P(t) = exp(-i t P / 2), every t_i an independent formal parameter.
Instance: n qubits and the M target Pauli strings (character i = qubit i, positive sign). Ranked instances also carry the
generator jw-hubbard-open-v1 and its L; the verifier regenerates the list and requires `paulis` to equal it. The debug
instance carries only n and paulis.
Artifact {"ops": [...]}, at most 50,000 operations: ["H",q] ["S",q] ["SDG",q] ["X",q] ["Z",q] ["CX",c,t] ["R",i,q,s].
R(i,q,s) is the physical gate exp(-i s t_i Z_q / 2), s = -1 or 1. The 2,097,152-byte file cap is applied by run_one.

Exact method (the spec's signed frame). A Pauli is (p, x, z), meaning i^p X^x Z^z with p mod 4 and integer bitsets x, z.
The frame keeps A_q = F^dag X_q F and B_q = F^dag Z_q F, where F is the product of the Clifford gates so far (rotations
omitted). A rotation R(i,q,s) placed after F equals F R_{s B_q}(t_i), so its pulled-back axis is s B_q. Accept iff every
label occurs exactly once, s B_q equals P_i (sign included) at its occurrence, every anticommuting predecessor P_j (j < i)
occurs before P_i, and every A_q, B_q is back at +X_q, +Z_q at the end. Then the circuit is F_final times the product of
the R_{P_i}(t_i) in an order that differs from the target order only by exchanges of commuting rotations, and F_final is
the identity up to a global phase. Score (cx, gates), both minimized, counted from the operation list.
Cost: O((G + M^2) n / 64) word operations; G <= 50,000 operations, M <= 512 labels, n <= 128 qubits.
"""
import os
import time

from common import Reject, exact_keys, is_int, load_instance, short, want_list

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

BOARD = "pauli-network"
GENERATOR = "jw-hubbard-open-v1"
HEAD_KEYS = ["board", "instance", "mode", "cpu_seconds", "mem_mb"]
RANKED_KEYS = HEAD_KEYS + ["generator", "L", "n", "paulis"]
DEBUG_KEYS = HEAD_KEYS + ["n", "paulis"]
N_MAX, M_MAX, OPS_MAX = 128, 512, 50000
ARITY = {"H": 2, "S": 2, "SDG": 2, "X": 2, "Z": 2, "CX": 3, "R": 4}
OP_H, OP_S, OP_SDG, OP_X, OP_Z, OP_CX, OP_R = range(7)
CODE = {"H": OP_H, "S": OP_S, "SDG": OP_SDG, "X": OP_X, "Z": OP_Z, "CX": OP_CX, "R": OP_R}


class Budget(Exception):
    pass


def hubbard_paulis(L):
    """jw-hubbard-open-v1: site u = y L + x, qubit 2u + s; open-boundary bonds sorted (u < v); per bond spin 0 then 1,
    X Z..Z X then Y Z..Z Y on p = 2u + s, q = 2v + s; then Z_{2u} Z_{2u+1} for every site in increasing order."""
    n = 2 * L * L
    bonds = sorted([(y * L + x, y * L + x + 1) for y in range(L) for x in range(L - 1)] +
                   [(y * L + x, (y + 1) * L + x) for y in range(L - 1) for x in range(L)])
    out = []
    for u, v in bonds:
        for s in (0, 1):
            p, q = 2 * u + s, 2 * v + s
            for a in "XY":
                out.append("I" * p + a + "Z" * (q - p - 1) + a + "I" * (n - q - 1))
    for u in range(L * L):
        out.append("I" * (2 * u) + "ZZ" + "I" * (n - 2 * u - 2))
    return out


def pauli_bits(s):
    """Positive Pauli string -> (p, x, z): x on X and Y, z on Z and Y, p = number of Y mod 4 (Y = i X Z)."""
    x = z = ny = 0
    for q, ch in enumerate(s):
        if ch in "XY":
            x |= 1 << q
        if ch in "ZY":
            z |= 1 << q
        ny += ch == "Y"
    return ny % 4, x, z


def pauli_text(p, x, z):
    """Readable signed Pauli, e.g. '-X3 Z4 Z5 Y6' (only the non-identity factors; at most n <= 128 of them)."""
    if p & 1 != popcount(x & z) & 1:
        return f"(non-Hermitian i^{p} X^{x:#x} Z^{z:#x})"
    sign = "+" if (p - popcount(x & z)) % 4 == 0 else "-"
    body = " ".join("IXZY"[(x >> q & 1) | (z >> q & 1) << 1] + str(q) for q in range(max(x, z).bit_length())
                    if (x | z) >> q & 1)
    return sign + (body or "I")


def instance_problem(I):
    """None if the instance file is well formed (and matches its generator), else a description of the defect."""
    if not isinstance(I, dict):
        return "instance is not an object"
    ranked = "generator" in I
    if sorted(I) != sorted(RANKED_KEYS if ranked else DEBUG_KEYS):
        return "instance keys differ from the specification"
    n, P = I["n"], I["paulis"]
    if I["board"] != BOARD or not is_int(n) or not 1 <= n <= N_MAX:
        return "instance header is malformed"
    if not isinstance(P, list) or not 1 <= len(P) <= M_MAX or \
            not all(isinstance(s, str) and len(s) == n and s.strip("IXYZ") == "" and s.strip("I") for s in P):
        return f"paulis must be 1..{M_MAX} non-identity strings of length n over I, X, Y, Z"
    if ranked:
        L = I["L"]
        if I["generator"] != GENERATOR or not is_int(L) or not 2 <= L or 2 * L * L != n:
            return "generator header is malformed"
        if P != hubbard_paulis(L):
            return "instance paulis differ from the regenerated jw-hubbard-open-v1 list"
    return None


def parse(art, n, M):
    """Strict artifact parsing -> list of tuples (code, a, b, c)."""
    exact_keys(art, ["ops"])
    ops = want_list(art["ops"], "ops", max_len=OPS_MAX)
    out = []
    for k, op in enumerate(ops):
        if not isinstance(op, list) or not op or not isinstance(op[0], str):
            raise Reject(f"ops[{k}] must be a list starting with an operation name")
        name = op[0]
        if name not in ARITY:
            raise Reject(f"ops[{k}]: unknown operation {short(name)} (allowed: H, S, SDG, X, Z, CX, R)")
        if len(op) != ARITY[name]:
            raise Reject(f"ops[{k}]: {name} takes {ARITY[name] - 1} integer operand(s), got {len(op) - 1}")
        for v in op[1:]:
            if not is_int(v):
                raise Reject(f"ops[{k}]: operands must be JSON integers, got {short(v)}")
        if name == "R":
            i, q, s = op[1], op[2], op[3]
            if not 0 <= i < M:
                raise Reject(f"ops[{k}]: rotation label {i} is outside 0..{M - 1}")
            if not 0 <= q < n:
                raise Reject(f"ops[{k}]: qubit {q} is outside 0..{n - 1}")
            if s not in (-1, 1):
                raise Reject(f"ops[{k}]: rotation sign must be -1 or 1, got {s}")
            out.append((OP_R, i, q, s))
        else:
            for q in op[1:]:
                if not 0 <= q < n:
                    raise Reject(f"ops[{k}]: qubit {q} is outside 0..{n - 1}")
            if name == "CX":
                if op[1] == op[2]:
                    raise Reject(f"ops[{k}]: CX control and target are both qubit {op[1]}")
                out.append((OP_CX, op[1], op[2], 0))
            else:
                out.append((CODE[name], op[1], 0, 0))
    return out


def predecessors(paulis, t0, budget):
    """pred[i] = bitset of the labels j < i whose target anticommutes with P_i; also the number of such pairs."""
    pred, pairs = [], 0
    for i, (_, xi, zi) in enumerate(paulis):
        if time.process_time() - t0 > budget:
            raise Budget
        m = 0
        for j in range(i):
            _, xj, zj = paulis[j]
            if popcount(xi & zj ^ zi & xj) & 1:
                m |= 1 << j
        pred.append(m)
        pairs += popcount(m)
    return pred, pairs


def replay(ops, n, paulis, pred, t0, budget):
    """Signed-frame replay; raises Reject at the first failing operation, returns None when everything holds."""
    ap, ax, az = [0] * n, [1 << q for q in range(n)], [0] * n   # A_q = i^ap X^ax Z^az
    bp, bx, bz = [0] * n, [0] * n, [1 << q for q in range(n)]   # B_q
    used = 0
    for k, (c, a, b, s) in enumerate(ops):
        if not k & 1023 and time.process_time() - t0 > budget:
            raise Budget
        if c == OP_CX:   # A_a <- A_a A_b, B_b <- B_a B_b
            ap[a] = (ap[a] + ap[b] + 2 * (popcount(az[a] & ax[b]) & 1)) & 3
            ax[a] ^= ax[b]
            az[a] ^= az[b]
            bp[b] = (bp[a] + bp[b] + 2 * (popcount(bz[a] & bx[b]) & 1)) & 3
            bx[b] ^= bx[a]
            bz[b] ^= bz[a]
        elif c == OP_R:
            i, q = a, b
            if used >> i & 1:
                raise Reject(f"ops[{k}]: rotation label {i} occurs a second time (each label exactly once)")
            got = ((bp[q] + (2 if s < 0 else 0)) & 3, bx[q], bz[q])
            if got != paulis[i]:
                raise Reject(f"ops[{k}]: R({i}, {q}, {s}) rotates about {pauli_text(*got)} but target {i} is "
                             f"{pauli_text(*paulis[i])}")
            miss = pred[i] & ~used
            if miss:
                j = (miss & -miss).bit_length() - 1
                raise Reject(f"ops[{k}]: rotation {i} comes before rotation {j}, but j < i and their targets "
                             "anticommute (they must keep the target order)")
            used |= 1 << i
        elif c == OP_H:
            ap[a], ax[a], az[a], bp[a], bx[a], bz[a] = bp[a], bx[a], bz[a], ap[a], ax[a], az[a]
        elif c == OP_S or c == OP_SDG:   # A <- -i A B (S) or +i A B (SDG)
            ap[a] = (ap[a] + bp[a] + 2 * (popcount(az[a] & bx[a]) & 1) + (3 if c == OP_S else 1)) & 3
            ax[a] ^= bx[a]
            az[a] ^= bz[a]
        elif c == OP_X:  # B <- -B
            bp[a] = (bp[a] + 2) & 3
        else:            # Z: A <- -A
            ap[a] = (ap[a] + 2) & 3
    M = len(paulis)
    if used != (1 << M) - 1:
        miss = ~used & ((1 << M) - 1)
        raise Reject(f"rotation label {(miss & -miss).bit_length() - 1} never occurs ({popcount(miss)} of {M} "
                     "labels missing)")
    for q in range(n):
        if (ap[q], ax[q], az[q]) != (0, 1 << q, 0) or (bp[q], bx[q], bz[q]) != (0, 0, 1 << q):
            raise Reject(f"the final Clifford frame is not the identity: X{q} is mapped to {pauli_text(ap[q], ax[q], az[q])} "
                         f"and Z{q} to {pauli_text(bp[q], bx[q], bz[q])} (signs count; undo every basis change and CX)")


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


def check(I, art, cpu_budget=None):
    """verify() on an instance given as a loaded dict (tests and the cost measurement pass instances this way)."""
    t0 = time.process_time()
    bad = instance_problem(I)
    if bad:
        return {"status": "error", "detail": {"reason": f"invalid instance file: {bad}"}}
    if cpu_budget is None:
        cpu_budget = 0.9 * float(os.environ.get("QARENA_CPU", I["cpu_seconds"]))
    n = I["n"]
    paulis = [pauli_bits(s) for s in I["paulis"]]
    M = len(paulis)
    ops = parse(art, n, M)
    try:
        pred, pairs = predecessors(paulis, t0, cpu_budget)
        replay(ops, n, paulis, pred, t0, cpu_budget)
        if time.process_time() - t0 > cpu_budget:
            raise Budget
    except Budget:
        return {"status": "error", "detail": {"reason": "verification did not finish within the CPU budget (not ranked)"}}
    cx = sum(1 for o in ops if o[0] == OP_CX)
    rot = sum(1 for o in ops if o[0] == OP_R)
    return {"status": "accepted", "score": {"cx": cx, "gates": len(ops)},
            "detail": {"n": n, "rotations": M, "cx": cx, "gates": len(ops), "single_qubit_cliffords": len(ops) - cx - rot,
                       "anticommuting_pairs": pairs,
                       "check": f"all {M} rotations occur once with the exact signed target axis, the {pairs} "
                                "anticommuting pairs keep the target order, and the final Clifford frame is the identity"}}