"""Board constant-add: exact modular constant addition with bounded clean workspace (SPEC_v3_final s.3).

Target |x> -> |x + c mod 2^n>, or for controlled instances |q, x> -> |q, x + q*c mod 2^n> (control = wire n).
Artifact {"ancillas": A, "ops": [...]}. Wires: data 0..n-1, control n (controlled only), then A workspace wires that start
live in |0>. Ops: ["X",t] ["CX",c,t] ["CCX",c1,c2,t] ["MX",q] ["RZ",q] ["PZ",q,bits,value].
MX measures a live workspace wire in the X basis (record bit 0 = +1 outcome) and deactivates it; RZ prepares an inactive
workspace wire in |0>; PZ applies Z to q exactly when (value XOR the XOR of the listed earlier records) is 1 (this is NOT the
"parity equals value" rule of ftprep/emitter). PZ may act on data, control and workspace wires.

Exact method: every live wire holds a function of the input only, kept as a packed Boolean column (Python int, bit x = the
wire value on input x, x over all 2^n or 2^(n+1) data/control inputs). X, CX, CCX update columns. An X measurement of a wire
holding f(x), with outcome r, multiplies the branch by (-1)^(r f(x)) / sqrt(2); PZ adds the current column of q to the phase
coefficients of the named records (and to the constant p_0 when value = 1). So the branch phase is
(-1)^(p_0(x) + sum_j r_j p_{j+1}(x)). The artifact is accepted iff the final data columns equal the target, the control is
unchanged, every workspace wire is live and zero, and every phase coefficient is a constant column (all 0 or all 1): then each
measurement branch is 2^(-M/2) times the target isometry up to a global phase (necessary as well as sufficient: a
non-constant coefficient j gives a wrong branch at r = e_j). Score (toffoli, gates, ancillas), all minimized.
Cost: O((G + M + R) 2^(n+1) / 64) word visits, G <= 20000 ops, M <= 512 records, R <= 8192 record references.
"""
import os
import time

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"))

ARITY = {"X": 2, "CX": 3, "CCX": 4, "MX": 2, "RZ": 2, "PZ": 4}
OPX, OPCX, OPCCX, OPMX, OPRZ, OPPZ = range(6)
CODE = {"X": OPX, "CX": OPCX, "CCX": OPCCX, "MX": OPMX, "RZ": OPRZ, "PZ": OPPZ}


class Budget(Exception):
    pass


def check_instance(I):
    """The instance file must be self-consistent; anything else is an internal inconsistency (status error)."""
    ok = (isinstance(I, dict) and all(is_int(I.get(k)) for k in ("n", "constant", "ancillas_max", "gates_max", "measurements_max",
                                                                  "phase_refs_max", "cpu_seconds"))
          and isinstance(I.get("controlled"), bool) and 1 <= I["n"] <= 20 and 0 <= I["constant"] < 1 << I["n"]
          and I["ancillas_max"] >= 0 and I["gates_max"] >= 0 and I["measurements_max"] >= 0 and I["phase_refs_max"] >= 0)
    if not ok:
        raise RuntimeError("instance file is inconsistent")


def parse_ops(ops, n, ctl, A, mmax, rmax):
    """Strict parse plus the static life cycle (it does not depend on the input values). Returns (parsed, live, nrec, refs)."""
    W = n + ctl + A
    live = [True] * W
    first_ws = n + ctl
    nrec = refs = 0
    out = []

    def wire(v, idx, what):
        return want_int(v, 0, W - 1, f"ops[{idx}] {what}")

    def need_live(w, idx):
        if not live[w]:
            raise Reject(f"ops[{idx}]: wire {w} is not live (it was measured and not re-prepared)")
        return w

    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 starting with an operation name")
        name = op[0]
        if name not in ARITY:
            raise Reject(f"ops[{idx}]: unsupported operation {short(name)}")
        if len(op) != ARITY[name]:
            raise Reject(f"ops[{idx}]: {name} takes {ARITY[name] - 1} operand(s), got {len(op) - 1}")
        if name in ("X", "CX", "CCX"):
            ws = [need_live(wire(v, idx, "operand"), idx) for v in op[1:]]
            if len(set(ws)) != len(ws):
                raise Reject(f"ops[{idx}]: {name} operands must be distinct")
            out.append((CODE[name], *ws))
        elif name == "MX":
            q = wire(op[1], idx, "operand")
            if q < first_ws:
                raise Reject(f"ops[{idx}]: MX may only measure a workspace wire, not data or control wire {q}")
            need_live(q, idx)
            live[q] = False
            nrec += 1
            if nrec > mmax:
                raise Reject(f"ops[{idx}]: more than measurements_max = {mmax} measurements")
            out.append((OPMX, q))
        elif name == "RZ":
            q = wire(op[1], idx, "operand")
            if q < first_ws:
                raise Reject(f"ops[{idx}]: RZ may only prepare a workspace wire, not data or control wire {q}")
            if live[q]:
                raise Reject(f"ops[{idx}]: RZ on live wire {q} (a live wire cannot be reset)")
            live[q] = True
            out.append((OPRZ, q))
        else:  # PZ
            q = need_live(wire(op[1], idx, "operand"), idx)
            bits = want_list(op[2], f"ops[{idx}] record list", max_len=nrec)
            prev = -1
            for b in bits:
                if not is_int(b) or not prev < b < nrec:
                    raise Reject(f"ops[{idx}]: record list must be strictly increasing indices of existing records "
                                 f"(0..{nrec - 1} at this point), got {short(op[2])}")
                prev = b
            refs += len(bits)
            if refs > rmax:
                raise Reject(f"ops[{idx}]: more than phase_refs_max = {rmax} record references in PZ operations")
            out.append((OPPZ, q, tuple(bits), want_int(op[3], 0, 1, f"ops[{idx}] value")))
    return out, live, nrec, refs


def input_columns(k):
    """Column of wire i (i < k): bit x is bit i of x, for x in 0..2^k-1. Returns (columns, all-ones mask)."""
    D = 1 << k
    mask = (1 << D) - 1
    cols = []
    for i in range(k):
        col, L = (((1 << (1 << i)) - 1) << (1 << i)), 1 << (i + 1)
        while L < D:
            col |= col << L
            L <<= 1
        cols.append(col)
    return cols, mask


def target_columns(n, c, ctl, init, mask):
    """Packed ripple-carry addition of the constant (c, or q*c for the control column q) to the data columns."""
    out, carry = [], 0
    for i in range(n):
        a = init[i]
        b = (init[n] if ctl else mask) if c >> i & 1 else 0
        t = a ^ b
        out.append(t ^ carry)
        carry = (a & b) | (carry & t)
    return out


def first_input(col):
    return (col & -col).bit_length() - 1


def describe(x, n, ctl):
    return f"x = {x & ((1 << n) - 1)}" + (f", control = {x >> n}" if ctl else "")


def verify(inst, art, time_budget=None):
    return check(load_instance("constadd", 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)
    n, c, ctl = I["n"], I["constant"], int(I["controlled"])
    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()
    exact_keys(art, ["ancillas", "ops"])
    A = want_int(art["ancillas"], 0, I["ancillas_max"], "ancillas")
    ops = want_list(art["ops"], "ops", max_len=I["gates_max"])
    parsed, live, nrec, refs = parse_ops(ops, n, ctl, A, I["measurements_max"], I["phase_refs_max"])

    k = n + ctl
    W = k + A
    init, mask = input_columns(k)
    val = init + [0] * A
    p0 = 0
    coef = []
    try:
        for i, op in enumerate(parsed):
            if not i & 255 and time.process_time() - t0 > budget:
                raise Budget
            t = op[0]
            if t == OPX:
                val[op[1]] ^= mask
            elif t == OPCX:
                val[op[2]] ^= val[op[1]]
            elif t == OPCCX:
                val[op[3]] ^= val[op[1]] & val[op[2]]
            elif t == OPMX:
                coef.append(val[op[1]])
                val[op[1]] = 0
            elif t == OPRZ:
                val[op[1]] = 0
            else:
                f = val[op[1]]
                if op[3]:
                    p0 ^= f
                for j in op[2]:
                    coef[j] ^= f
        want = target_columns(n, c, ctl, init, mask)
    except Budget:
        return {"status": "error", "detail": {"reason": "verification did not finish within the CPU budget (not ranked)"}}

    for i in range(n):
        if val[i] != want[i]:
            diff = val[i] ^ want[i]
            raise Reject(f"data wire {i} ends with a wrong value on {popcount(diff)} of {1 << k} inputs "
                         f"(first: {describe(first_input(diff), n, ctl)})")
    if ctl and val[n] != init[n]:
        diff = val[n] ^ init[n]
        raise Reject(f"the control wire is changed on {popcount(diff)} of {1 << k} inputs "
                     f"(first: {describe(first_input(diff), n, ctl)})")
    for w in range(k, W):
        if not live[w]:
            raise Reject(f"workspace wire {w} is not live at the end (measured and not re-prepared)")
        if val[w]:
            raise Reject(f"workspace wire {w} is not clean: nonzero on {popcount(val[w])} of {1 << k} inputs "
                         f"(first: {describe(first_input(val[w]), n, ctl)})")
    for j, col in enumerate([p0] + coef):
        if col not in (0, mask):
            ones = popcount(col)
            who = "the unconditional / constant phase" if j == 0 else f"the phase attached to measurement record {j - 1}"
            raise Reject(f"{who} depends on the input ({ones} of {1 << k} inputs have phase flip 1; "
                         f"first input with a different value than input 0: "
                         f"{describe(first_input(col ^ (mask if col & 1 else 0)), n, ctl)}): the measured value is not "
                         "corrected, so some measurement outcomes give a wrong or entangled state")

    toffoli = sum(1 for op in parsed if op[0] == OPCCX)
    return {"status": "accepted", "score": {"toffoli": toffoli, "gates": len(ops), "ancillas": A},
            "detail": {"n": n, "controlled": bool(ctl), "toffoli": toffoli, "gates": len(ops), "ancillas": A,
                       "measurements": nrec, "phase_refs": refs,
                       "check": f"all {1 << k} inputs: target output, unchanged control, clean workspace, "
                                f"{nrec + 1} constant phase coefficients"}}