"""Board ecdlp-pointadd: exact coherent point addition on y^2 = x^3 + 7 over F_p (SPEC_v5 s2 with review_v5_fable P1-P3, P5, P6
and the chair's P2 cell).

Target: R = P + Q by the chord formula lambda = (y_Q - y_P) / (x_Q - x_P), x_R = lambda^2 - x_P - x_Q,
y_R = lambda (x_P - x_R) - y_P (mod p), written into the P register, with
  kind offset  (Q classical, the instance's "offset"): wires P.x = 0..b-1, P.y = b..2b-1, then workspace;
  kind quantum (Q read-only quantum addend): P as above, Q.x = 2b..3b-1, Q.y = 3b..4b-1, then workspace.
Coordinates are little-endian, b = p.bit_length(). The promised domain (every input is checked, no sampling) is
  domain curve, offset:  every finite curve point P (the instance's "points") with x_P != x_Q and x_R != x_Q;
  domain curve, quantum: every ordered pair (P, Q) of finite curve points with x_P != x_Q and x_R != x_Q;
  domain plane, offset:  every (x, y) in [0, p)^2 with x != x_Q and x_R != x_Q (the chord map; no curve membership).
Artifact {"ancillas": A, "ops": [...]} with the constant-add vocabulary ["X",t] ["CX",c,t] ["CCX",c1,c2,t] ["MX",q] ["RZ",q]
["PZ",q,bits,value]; Q wires may be controls and PZ targets but never targets of X, CX or CCX; MX measures a live workspace wire in
X (record 0 = outcome +1) and deactivates it; RZ prepares an inactive workspace wire in |0>; PZ applies Z_q^(value XOR r_bits).

Exact method (packed Boolean columns, as in constadd.py): the promised inputs are processed in chunks of at most chunk_inputs
(offset inputs in the order of "points", plane inputs ordered by (x, y), quantum inputs ordered by Q, then P). In a chunk every
live wire holds a Python int whose bit k is the wire's value on input k; X XORs the chunk's own all-ones mask (exactly its input
count, review P3). The branch phase is (-1)^(p_0(u) + sum_j r_j p_{j+1}(u)): MX(q) creates a coefficient equal to q's column, PZ
XORs q's column into the coefficients of the listed records (and into p_0 when value = 1). A chunk passes iff the P wires equal
the target columns, the Q wires are unchanged, every workspace wire is live and zero, and every coefficient is 0 or the chunk mask;
the constant (0 or 1) of each coefficient is fixed by the first chunk and every later chunk must give the same constant (P3: the
comparison is of the constant, not of the packed integer). Then every measurement branch is 2^(-M/2) times the target isometry
up to a branch sign that does not depend on the input; the target is injective on the promised domain, so the full output check
also rules out loss of input information. Score (toffoli, gates, ancillas), minimized lexicographically.
Cost: O((G + M + R) * N / 64) word operations for G ops, M records, R record references and N promised inputs.
"""
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}

PRIMES = (11, 251, 4093)
OFFSETS = {11: [2, 2], 251: [0, 42], 4093: [5, 65]}
CEILINGS = {"qubits_max": 512, "gates_max": 1048576, "measurements_max": 4096, "phase_refs_max": 65536,
            "chunk_inputs": 262144, "cpu_seconds": 900, "mem_mb": 2048}
INST_KEYS = ("board", "instance", "mode", "kind", "domain", "p", "a", "b", "coordinate_bits", "points", "offset",
             "promised_inputs") + tuple(CEILINGS)


class Budget(Exception):
    pass


def curve_points(p):
    """All finite points of y^2 = x^3 + 7 over F_p, lexicographic."""
    roots = {}
    for y in range(p):
        roots.setdefault(y * y % p, []).append(y)
    return [[x, y] for x in range(p) for y in roots.get((x * x * x + 7) % p, [])]


def chord(p, P, Q):
    """P + Q by the chord formula (needs x_P != x_Q; curve membership is not used)."""
    lam = (Q[1] - P[1]) * pow(Q[0] - P[0], -1, p) % p
    xr = (lam * lam - P[0] - Q[0]) % p
    return [xr, (lam * (P[0] - xr) - P[1]) % p]


def check_instance(I):
    """The instance file must be self-consistent; anything else is an internal inconsistency (status error)."""
    ok = isinstance(I, dict) and set(I) == set(INST_KEYS)
    ok = ok and I["board"] == "ecdlp-pointadd" and I["kind"] in ("offset", "quantum") and I["domain"] in ("curve", "plane")
    ok = ok and all(is_int(I[k]) and 1 <= I[k] <= v for k, v in CEILINGS.items())
    ok = ok and is_int(I["p"]) and I["p"] in PRIMES and I["a"] == 0 and I["b"] == 7 and is_int(I["a"]) and is_int(I["b"])
    ok = ok and I["coordinate_bits"] == I["p"].bit_length() and is_int(I["promised_inputs"]) and I["promised_inputs"] > 0
    ok = ok and I["mode"] == ("ranked" if I["p"] == 4093 else "debug")
    if ok:
        p = I["p"]
        if I["domain"] == "curve":
            ok = I["points"] == curve_points(p)
        else:
            ok = I["points"] is None and I["kind"] == "offset"
        if I["kind"] == "offset":
            ok = ok and I["offset"] == OFFSETS[p]
        else:
            ok = ok and I["offset"] is None
        name = f"ec{p}-{I['kind']}" + ("-plane" if I["domain"] == "plane" else "")
        ok = ok and I["instance"] == name
    if not ok:
        raise RuntimeError("instance file is inconsistent")


def parse_ops(ops, ndata, nq, A, mmax, rmax):
    """Strict parse plus the static life cycle (it does not depend on input values). Returns (parsed, live, nrec, refs).
    Wires 0..ndata-1 hold P (and Q), of which the last nq (Q.x, Q.y) are read-only; workspace starts at ndata."""
    W = ndata + A
    live = [True] * W
    first_q = ndata - nq
    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")
            if first_q <= ws[-1] < ndata:
                raise Reject(f"ops[{idx}]: {name} targets wire {ws[-1]} of the read-only quantum addend Q "
                             "(Q wires may only be controls or PZ targets)")
            out.append((CODE[name], *ws))
        elif name == "MX":
            q = wire(op[1], idx, "operand")
            if q < ndata:
                raise Reject(f"ops[{idx}]: MX may only measure a workspace wire, not data 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 < ndata:
                raise Reject(f"ops[{idx}]: RZ may only prepare a workspace wire, not data 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 domain_blocks(I):
    """Yield the promised inputs in verifier order as numpy blocks (xp, yp, xq, yq, xr, yr) of int64 arrays."""
    import numpy as np
    p = I["p"]
    inv = np.zeros(p, dtype=np.int64)
    inv[1:] = [pow(i, -1, p) for i in range(1, p)]

    def add(xp, yp, xq, yq):
        """Chord formula, vectorized; keeps only x_P != x_Q and x_R != x_Q (review P1)."""
        keep = xp != xq
        xp, yp, xq, yq = xp[keep], yp[keep], xq[keep], yq[keep]
        lam = (yq - yp) % p * inv[(xq - xp) % p] % p
        xr = (lam * lam - xp - xq) % p
        yr = (lam * (xp - xr) - yp) % p
        keep = xr != xq
        return xp[keep], yp[keep], xq[keep], yq[keep], xr[keep], yr[keep]

    if I["domain"] == "plane":
        xq, yq = I["offset"]
        ys = np.arange(p, dtype=np.int64)
        for x in range(p):
            n = len(ys)
            yield add(np.full(n, x, dtype=np.int64), ys, np.full(n, xq, dtype=np.int64), np.full(n, yq, dtype=np.int64))
    else:
        pts = np.array(I["points"], dtype=np.int64)
        n = len(pts)
        if I["kind"] == "offset":
            xq, yq = I["offset"]
            yield add(pts[:, 0], pts[:, 1], np.full(n, xq, dtype=np.int64), np.full(n, yq, dtype=np.int64))
        else:
            for xq, yq in I["points"]:
                yield add(pts[:, 0], pts[:, 1], np.full(n, xq, dtype=np.int64), np.full(n, yq, dtype=np.int64))


def domain_chunks(I, B):
    """Concatenate domain_blocks into chunks of exactly B inputs (the last one may be shorter)."""
    import numpy as np
    buf, have = [], 0
    for blk in domain_blocks(I):
        buf.append(blk)
        have += len(blk[0])
        while have >= B:
            cat = [np.concatenate([b[i] for b in buf]) for i in range(6)]
            yield [c[:B] for c in cat]
            rest = [c[B:] for c in cat]
            buf, have = [rest], len(rest[0])
    if have:
        yield [np.concatenate([b[i] for b in buf]) for i in range(6)]


def pack(v, b):
    """b packed columns of the little-endian bits of the int64 array v (bit k of column i = bit i of v[k])."""
    import numpy as np
    return [int.from_bytes(np.packbits(((v >> i) & 1).astype(np.uint8), bitorder="little").tobytes(), "little")
            for i in range(b)]


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


def describe(ch, k, quantum):
    s = f"P = ({int(ch[0][k])}, {int(ch[1][k])})"
    return s + (f", Q = ({int(ch[2][k])}, {int(ch[3][k])})" if quantum else "")


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


def check(I, art, time_budget=None, abort_before_chunk=None):
    """verify() on an instance given as a loaded dict. abort_before_chunk (tests only) simulates an interruption."""
    check_instance(I)
    p, b, quantum = I["p"], I["coordinate_bits"], I["kind"] == "quantum"
    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"])
    ndata = 4 * b if quantum else 2 * b
    nq = 2 * b if quantum else 0
    A = want_int(art["ancillas"], 0, I["qubits_max"] - ndata, "ancillas")
    ops = want_list(art["ops"], "ops", max_len=I["gates_max"])
    parsed, live, nrec, refs = parse_ops(ops, ndata, nq, A, I["measurements_max"], I["phase_refs_max"])
    W = ndata + A
    for w in range(ndata, W):
        if not live[w]:
            raise Reject(f"workspace wire {w} is not live at the end (measured and not re-prepared)")

    const = None  # constants (0/1) of p_0, p_1, ..., p_M, fixed by the first chunk
    total = nchunks = 0
    try:
        for ch in domain_chunks(I, I["chunk_inputs"]):
            if abort_before_chunk is not None and nchunks == abort_before_chunk:
                raise Budget
            if time.process_time() - t0 > budget:
                raise Budget
            n = len(ch[0])
            mask = (1 << n) - 1
            init = pack(ch[0], b) + pack(ch[1], b) + (pack(ch[2], b) + pack(ch[3], b) if quantum else [])
            val = init + [0] * A
            p0 = 0
            coef = []
            for i, op in enumerate(parsed):
                if not i & 1023 and time.process_time() - t0 > budget:
                    raise Budget
                t = op[0]
                if t == OPCX:
                    val[op[2]] ^= val[op[1]]
                elif t == OPCCX:
                    val[op[3]] ^= val[op[1]] & val[op[2]]
                elif t == OPX:
                    val[op[1]] ^= mask
                elif t == OPPZ:
                    f = val[op[1]]
                    if op[3]:
                        p0 ^= f
                    for j in op[2]:
                        coef[j] ^= f
                elif t == OPMX:
                    coef.append(val[op[1]])
                    val[op[1]] = 0
                else:
                    val[op[1]] = 0
            want = pack(ch[4], b) + pack(ch[5], b)
            for i in range(2 * b):
                if val[i] != want[i]:
                    k = first_input(val[i] ^ want[i])
                    reg = f"P.{'xy'[i // b]} bit {i % b} (wire {i})"
                    raise Reject(f"{reg} ends with a wrong value on {popcount(val[i] ^ want[i])} inputs of a chunk "
                                 f"of {n} (first: {describe(ch, k, quantum)}; expected P + Q = "
                                 f"({int(ch[4][k])}, {int(ch[5][k])}))")
            for i in range(2 * b, ndata):
                if val[i] != init[i]:  # unreachable: Q wires are never targets; kept as a guard
                    raise Reject(f"Q wire {i} is changed (first: {describe(ch, first_input(val[i] ^ init[i]), quantum)})")
            for w in range(ndata, W):
                if val[w]:
                    raise Reject(f"workspace wire {w} is not clean: nonzero on {popcount(val[w])} inputs of a chunk of "
                                 f"{n} (first: {describe(ch, first_input(val[w]), quantum)})")
            cols = [p0] + coef
            if const is None:
                const = [col & 1 for col in cols]
            for j, col in enumerate(cols):
                ref = mask if const[j] else 0
                if col != ref:
                    k = first_input(col ^ ref)
                    who = ("the unconditional phase p_0" if j == 0 else
                           f"the phase coefficient of measurement record {j - 1}")
                    raise Reject(f"{who} depends on the input: it is {const[j]} on the first promised input and "
                                 f"{1 - const[j]} on {describe(ch, k, quantum)}; the measurement is not corrected, so some "
                                 "outcomes give a wrong or entangled state")
            total += n
            nchunks += 1
            del val, coef, cols, init, want
    except Budget:
        return {"status": "error", "detail": {"reason": "verification did not finish within the CPU budget (not ranked)",
                                              "chunks_done": nchunks}}
    if total != I["promised_inputs"]:
        raise RuntimeError(f"promised domain has {total} inputs, instance says {I['promised_inputs']}")

    toffoli = sum(1 for op in parsed if op[0] == OPCCX)
    return {"status": "accepted", "score": {"toffoli": toffoli, "gates": len(ops), "ancillas": A},
            "detail": {"kind": I["kind"], "domain": I["domain"], "p": p, "toffoli": toffoli, "gates": len(ops),
                       "ancillas": A, "measurements": nrec, "phase_refs": refs, "chunks": nchunks,
                       "check": f"all {total} promised inputs: P + Q in the P register, "
                                + ("unchanged Q, " if quantum else "")
                                + f"clean live workspace, {nrec + 1} phase coefficients constant on the whole domain"}}