"""Board qrom: exact coherent table lookup with a clean workspace (SPEC_v3_final s2).

Target: the promised-input isometry |x>|0^b>|0^s> -> |x>|D[x]>|0^s>. Wires: address 0..a-1 (wire i = weight 2^i of x),
output a..a+b-1 (wire a+j = weight 2^j of the word D[x]), declared workspace a+b..a+b+s-1.
Artifact {"ancillas": s, "ops": [...]}, 0 <= s <= ancillas_max, at most gates_max ops, each exactly ["X", t],
["CX", c, t] or ["CCX", c1, c2, t] with JSON-integer wire indices below a+b+s, distinct within the op.
Accept iff for every address x the circuit maps (address x, output 0, workspace 0) to (x, D[x], 0). X, CX and CCX are
permutation matrices with entries +1, so this basis check over all 2^a addresses proves the coherent map, phases
included. Score (toffoli, gates, ancillas) = (number of CCX, len(ops), s), all minimized; the submitted operations are
counted as written.

Exact method: bit slicing. Wire w holds one Python int over all 2^a addresses (bit x = the value of w on address x),
so X, CX and CCX are one XOR (and one AND) on whole columns; at the end the columns are compared with the address and
table columns. Work O(len(ops) * 2^a / 64) machine words; the CPU budget is checked every CHUNK operations.
verify() first checks the stored instance against its generator (SPEC_v3_final s2; family hashpair for pair16-b16-wide
since the CEO's decision of 2026-11-14); a mismatch is an instance error.
"""
import hashlib
import os
import time

from common import Reject, exact_keys, want_int, want_list, 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}
CHUNK = 1 << 14
LIST_MAX = 8   # wires named in a rejection message


INSTANCE_KEYS = {"board", "instance", "mode", "cpu_seconds", "mem_mb", "family", "address_bits", "output_bits",
                 "ancillas_max", "gates_max", "seed", "data"}


def pair_table(a, b):
    """The spec's pair recipe: (seed, table), address u + S*v holding the quantized pair coefficient of sites u, v."""
    L = 1 << (a // 4)
    S = L * L
    seed = state = 0x51415245 ^ a ^ (b << 8)
    charge = []
    for _ in range(S):
        state = (1664525 * state + 1013904223) & 0xffffffff
        charge.append(1 + state % 31)
    top = (1 << b) - 1
    return seed, [top * charge[u] * charge[v] // (961 * (1 + (u % L - v % L) ** 2 + (u // L - v // L) ** 2))
                  for v in range(S) for u in range(S)]


def hashpair_table(a, b):
    """Family hashpair (pair16-b16-wide): sites u, v < S = 2^(a/2), address u + S*v. For u <= v the word is SHA-256 of
    the ASCII text "qrom <seed> <u> <v>", its first ceil(b/8) bytes read little endian and cut to b bits; the table is
    symmetric. The seed follows the pair recipe's formula. Returns (seed, table)."""
    S = 1 << (a // 2)
    seed = 0x51415245 ^ a ^ (b << 8)
    nbytes, mask = (b + 7) // 8, (1 << b) - 1
    table = [0] * (S * S)
    for u in range(S):
        for v in range(u, S):
            word = int.from_bytes(hashlib.sha256(f"qrom {seed} {u} {v}".encode("ascii")).digest()[:nbytes], "little")
            table[u + S * v] = table[v + S * u] = word & mask
    return seed, table


def validate_instance(I):
    """A stored instance must equal its generator output: the pair recipe, the hashpair recipe, or for family gray seed 0
    and the Gray code x ^ (x >> 1), which is the spec's table [0, 1, 3, 2] for a = 2. Any mismatch raises RuntimeError
    (status error)."""
    if type(I) is not dict or set(I) != INSTANCE_KEYS or I["board"] != "qrom":
        raise RuntimeError("invalid qrom instance: wrong fields")
    a, b, _, _, D = check_instance(I)
    if I["family"] == "pair" and a % 4 == 0:
        want = pair_table(a, b)
    elif I["family"] == "hashpair" and a % 2 == 0:
        want = hashpair_table(a, b)
    elif I["family"] == "gray":
        want = 0, [x ^ x >> 1 for x in range(1 << a)]
    else:
        raise RuntimeError(f"invalid qrom instance: unknown family {short(I['family'])}")
    if (I["seed"], D) != want:
        raise RuntimeError("invalid qrom instance: stored seed or table differs from its generator")


def check_instance(I):
    a, b, smax, gmax, D = I["address_bits"], I["output_bits"], I["ancillas_max"], I["gates_max"], I["data"]
    if not all(type(v) is int for v in (a, b, smax, gmax)) or not (1 <= a <= 24 and 1 <= b <= 64 and smax >= 0
                                                                   and gmax >= 0):
        raise RuntimeError("invalid qrom instance parameters")
    top = 1 << b
    if type(D) is not list or len(D) != 1 << a or not all(type(v) is int and 0 <= v < top for v in D):
        raise RuntimeError("invalid qrom instance table")
    return a, b, smax, gmax, D


def address_columns(a):
    """Column of address wire i: bit x is bit i of x (2^i zeros, 2^i ones, repeated)."""
    N, cols = 1 << a, []
    for i in range(a):
        col, width = ((1 << (1 << i)) - 1) << (1 << i), 2 << i
        while width < N:
            col |= col << width
            width <<= 1
        cols.append(col)
    return cols


def table_columns(D, b):
    """Column of output bit j: bit x is bit j of D[x]."""
    return [int("".join("1" if d >> j & 1 else "0" for d in reversed(D)), 2) for j in range(b)]


def parse(art, a, b, smax, gmax):
    exact_keys(art, ["ancillas", "ops"])
    s = want_int(art["ancillas"], 0, smax, "ancillas")
    ops = want_list(art["ops"], "ops", max_len=gmax)
    W = a + b + s
    toffoli = cx = 0
    for k, op in enumerate(ops):
        if type(op) is not list or not op:
            raise Reject(f"ops[{k}] must be a non-empty list such as [\"CX\", c, t], got {short(op)}")
        g = op[0]
        n = ARITY.get(g) if type(g) is str else None
        if n is None:
            raise Reject(f"ops[{k}]: unsupported operation {short(g)}; the only operations are X, CX and CCX "
                         f"(no phase gates, measurements or macros)")
        if len(op) != n:
            raise Reject(f"ops[{k}]: {g} takes exactly {n - 1} wire operand{'s' if n > 2 else ''}, got {len(op) - 1}")
        for w in op[1:]:
            if type(w) is not int or not 0 <= w < W:
                raise Reject(f"ops[{k}]: operand {short(w)} is not a wire index; wires are the JSON integers 0..{W - 1} "
                             f"({a} address, {b} output, {s} workspace)")
        if n == 3:
            if op[1] == op[2]:
                raise Reject(f"ops[{k}]: repeated operand in {short(op)}")
            cx += 1
        elif n == 4:
            if op[1] == op[2] or op[1] == op[3] or op[2] == op[3]:
                raise Reject(f"ops[{k}]: repeated operand in {short(op)}")
            toffoli += 1
    return s, ops, toffoli, cx


def wires_text(ws):
    return ", ".join(map(str, ws[:LIST_MAX])) + (f" and {len(ws) - LIST_MAX} more" if len(ws) > LIST_MAX else "")


def check(I, art, time_budget=None):
    """verify() on an instance given as a loaded dict (tests pass small ad hoc tables this way)."""
    a, b, smax, gmax, D = check_instance(I)
    budget = time_budget if time_budget is not None else float(os.environ.get("QARENA_CPU", I["cpu_seconds"])) - 30
    t0 = time.process_time()
    s, ops, toffoli, cx = parse(art, a, b, smax, gmax)
    N, W = 1 << a, a + b + s
    full = (1 << N) - 1
    addr = address_columns(a)
    col = addr + [0] * (b + s)
    for start in range(0, len(ops), CHUNK):
        if time.process_time() - t0 > budget:
            return {"status": "error", "detail": {"reason": "simulation did not finish within the CPU limit (not ranked)",
                                                  "ops_done": start, "ops": len(ops)}}
        for op in ops[start:start + CHUNK]:
            g = op[0]
            if g == "CCX":
                col[op[3]] ^= col[op[1]] & col[op[2]]
            elif g == "CX":
                col[op[2]] ^= col[op[1]]
            else:
                col[op[1]] ^= full

    tgt = table_columns(D, b)
    bad_addr = bad_out = bad_ws = 0
    for i in range(a):
        bad_addr |= col[i] ^ addr[i]
    for j in range(b):
        bad_out |= col[a + j] ^ tgt[j]
    for w in range(a + b, W):
        bad_ws |= col[w]
    bad = bad_addr | bad_out | bad_ws
    if bad:
        x = (bad & -bad).bit_length() - 1
        got = sum((col[a + j] >> x & 1) << j for j in range(b))
        what = []
        if bad_out >> x & 1:
            what.append(f"output word {got} instead of D[{x}] = {D[x]} (wrong output wires "
                        f"{wires_text([a + j for j in range(b) if (got ^ D[x]) >> j & 1])})")
        if bad_addr >> x & 1:
            what.append(f"address wires {wires_text([i for i in range(a) if (col[i] ^ addr[i]) >> x & 1])} changed")
        if bad_ws >> x & 1:
            what.append(f"workspace wires {wires_text([w for w in range(a + b, W) if col[w] >> x & 1])} not returned to 0")
        raise Reject(f"wrong on {popcount(bad)} of {N} addresses (output wrong on {popcount(bad_out)}, address changed "
                     f"on {popcount(bad_addr)}, workspace not clean on {popcount(bad_ws)}); first failing address {x}: "
                     + "; ".join(what))
    G = len(ops)
    return {"status": "accepted", "score": {"toffoli": toffoli, "gates": G, "ancillas": s},
            "detail": {"address_bits": a, "output_bits": b, "ancillas": s, "wires": W, "toffoli": toffoli, "cx": cx,
                       "x": G - toffoli - cx, "gates": G,
                       "check": f"all {N} addresses: address restored, output equals the table word, workspace back to 0",
                       "verify_seconds": round(time.process_time() - t0, 2)}}


def verify(inst, art, time_budget=None):
    I = load_instance("qrom", inst)
    validate_instance(I)
    return check(I, art, time_budget)