"""Board distillation-code: one-round T-state distillation by a triorthogonal matrix (SPEC_v5 s.3, review D1-D3).

Artifact {"n": n, "logical": [k strings], "stabilizers": [1..rank_max-k strings]}; each string has n characters '0'/'1'
(character j = column j). G = logical rows g_0..g_{k-1} followed by stabilizer rows; r = number of rows. Review D3
removed the submitted correction: the verifier derives the canonical one and returns it in detail.

Protocol (row-coordinate form of Bravyi-Haah): r row qubits start in |+>, one noisy T rotation is injected on each of
the n column parities, the Clifford correction C = prod_i S_i^{s_i} prod_{[i,j] in cz} CZ_ij is applied, and the
stabilizer row qubits are measured in X; the run is accepted only if every one of them gives +. Each input is
independently rho_p = (1-p)|T><T| + p Z|T><T|Z. An input error pattern e (one bit per column) is accepted iff G_0 e = 0
(G_0 = stabilizer rows) and then flips the logical outputs G_1 e (G_1 = logical rows).

Acceptance, in this order (each failure is a rejection with the reason):
  1. G has full row rank over GF(2) (r <= rank_max);
  2. logical rows have odd weight, stabilizer rows even weight;
  3. every pair of distinct rows has even overlap;
  4. every triple of distinct rows has even overlap;
  5. d_Z >= 3: no column has a zero stabilizer syndrome and a nonzero logical part (weight-1 error), and no two
     columns have the same stabilizer syndrome but different logical parts (weight-2 error);
  6. the exact conditional probability of any output error, q_bad = 1 - B/A, is at most error_num/error_den.
Given 2-4, s_i = (delta_i - |g_i|)/2 mod 4 (delta_i = 1 for logical rows) and cz = {[i,j] : |g_i & g_j|/2 odd}
cancel every residual phase: the corrected phase of row basis state x is zeta_8^(x_0 + ... + x_{k-1}), so the k logical
row qubits leave in T|+>.

Exact probabilities (the Fourier form of Pr(G_0 e = 0) and Pr(G e = 0)): with z = 1 - 2p and s = r - k,
  A = Pr(accept) = 2^-s W_0(z),  B = Pr(accept and no output error) = 2^-r W(z),
where W_0 and W are the weight enumerators of rowspan G_0 and rowspan G. With p = a/b, z = (b - 2a)/b and
  N_0 = sum_w W_0[w] (b-2a)^w b^(n-w),  N = sum_w W[w] (b-2a)^w b^(n-w)  (integers),
A = N_0 / (2^s b^n), B/A = N / (2^k N_0), q_bad = (2^k N_0 - N) / (2^k N_0), and cost = n / (k A) = n 2^s b^n / (k N_0),
the input T states consumed per accepted output T state (k comes from the instance: 8 on the ranked instances, 7 on the
debug instances).
Score, all minimized: cost, q_bad (exact reduced fractions in lowercase hex), n, rank r.

Weight enumerators: the first m = min(r', 12) rows of a row set (r' rows) are expanded once into the 2^m words of
their span; the other rows are visited in Gray-code order (one XOR per step), and for each such word h the 2^m words
h ^ l are counted by weight. Work 2^r' popcounts per enumerator (at most 2^24 + 2^17 = 16,908,288 at r = 24), memory
2^m words plus the two histograms. The CPU budget is polled between Gray steps and before the final comparison; an
exhausted budget returns status error, never a partial histogram.
"""
import os
import time
from collections import Counter
from fractions import Fraction

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

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

INSTANCE_KEYS = ("board", "instance", "mode", "n_min", "n_max", "k", "rank_max", "distance_min", "p_num", "p_den",
                 "error_num", "error_den", "cpu_seconds", "mem_mb")
LOW_ROWS = 12


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 set(I) == set(INSTANCE_KEYS) and I["board"] == "distillation-code"
          and isinstance(I["instance"], str) and I["mode"] in ("ranked", "debug")
          and all(is_int(I[k]) for k in INSTANCE_KEYS[3:])
          and 1 <= I["n_min"] <= I["n_max"] <= 4096 and 1 <= I["k"] < I["rank_max"] <= 26 and I["distance_min"] == 3
          and 0 <= I["p_num"] and 1 <= I["p_den"] and 2 * I["p_num"] <= I["p_den"]
          and 0 <= I["error_num"] <= I["error_den"] and 1 <= I["error_den"]
          and I["cpu_seconds"] >= 1 and I["mem_mb"] >= 1)
    if not ok:
        raise RuntimeError("instance file is inconsistent")


def parse(art, I):
    """Validate the artifact against instance I. Returns (n, logical rows, stabilizer rows) as int bitsets."""
    exact_keys(art, ["n", "logical", "stabilizers"])
    k, rmax = I["k"], I["rank_max"]
    n = want_int(art["n"], I["n_min"], I["n_max"], f"n (instance {I['instance']} allows {I['n_min']} <= n <= {I['n_max']})")
    lg = want_list(art["logical"], f"logical (exactly k = {k} rows)", max_len=k, min_len=k)
    st = want_list(art["stabilizers"], f"stabilizers (1 to rank_max - k = {rmax - k} rows)", max_len=rmax - k, min_len=1)
    L = [bitstring(x, n, f"logical[{i}]") for i, x in enumerate(lg)]
    S = [bitstring(x, n, f"stabilizers[{i}]") for i, x in enumerate(st)]
    return n, L, S


def row_name(i, k):
    return f"logical[{i}]" if i < k else f"stabilizers[{i - k}]"


def first_bits(v, cap=12):
    out, j = [], 0
    while v and len(out) < cap:
        if v & 1:
            out.append(j)
        v >>= 1
        j += 1
    return out


def check_rank(G, k):
    """Full row rank; on failure name a nonempty set of rows whose XOR is zero."""
    basis = []                                   # (pivot bit, reduced row, set of original rows combined)
    for i, g in enumerate(G):
        v, who = g, 1 << i
        for p, b, w in basis:
            if v >> p & 1:
                v ^= b
                who ^= w
        if v == 0:
            rows = [row_name(j, k) for j in first_bits(who, 24)]
            raise Reject(f"the rows are linearly dependent: the XOR of {', '.join(rows)} is zero (G must have full "
                         "row rank)")
        basis.append((v.bit_length() - 1, v, who))


def check_triorthogonal(G, k):
    r = len(G)
    for i, g in enumerate(G):
        w = popcount(g)
        if w % 2 != (1 if i < k else 0):
            want = "odd (logical rows)" if i < k else "even (stabilizer rows)"
            raise Reject(f"{row_name(i, k)} has weight {w}; it must be {want}")
    for i in range(r):
        for j in range(i + 1, r):
            o = popcount(G[i] & G[j])
            if o % 2:
                raise Reject(f"rows {row_name(i, k)} and {row_name(j, k)} overlap in {o} columns (odd); every pair of "
                             "distinct rows must have even overlap")
    for i in range(r):
        for j in range(i + 1, r):
            gij = G[i] & G[j]
            for l in range(j + 1, r):
                o = popcount(gij & G[l])
                if o % 2:
                    raise Reject(f"rows {row_name(i, k)}, {row_name(j, k)} and {row_name(l, k)} overlap in {o} columns "
                                 "(odd); every triple of distinct rows must have even overlap")


def columns(n, rows):
    """Column j as an int whose bit i is row i at column j."""
    cols = [0] * n
    for i, g in enumerate(rows):
        bit = 1 << i
        j = 0
        while g:
            if g & 1:
                cols[j] |= bit
            g >>= 1
            j += 1
    return cols


def check_distance(n, L, S):
    """d_Z >= 3 by every weight-1 and weight-2 error, compared through column syndromes. Returns (lam, sig) lists."""
    lam, sig = columns(n, L), columns(n, S)
    seen = {}
    for j in range(n):
        if sig[j] == 0 and lam[j]:
            raise Reject(f"d_Z = 1: column {j} has no stabilizer row on it but lies in a logical row, so a single input "
                         "error there is accepted and flips an output")
        if sig[j] in seen and lam[seen[sig[j]]] != lam[j]:
            i = seen[sig[j]]
            raise Reject(f"d_Z <= 2: columns {i} and {j} have the same stabilizer syndrome but different logical parts, "
                         f"so errors on both are accepted and flip outputs {first_bits(lam[i] ^ lam[j])}")
        seen.setdefault(sig[j], j)
    return lam, sig


def weight3_logical(n, lam, sig):
    """Number of weight-3 input errors that are accepted and flip at least one output (display only)."""
    by = {}
    for j in range(n):
        by.setdefault(sig[j], []).append(j)
    cnt = 0
    for a in range(n):
        for b in range(a + 1, n):
            for c in by.get(sig[a] ^ sig[b], ()):
                if c > b and lam[a] ^ lam[b] ^ lam[c]:
                    cnt += 1
    return cnt


def enumerate_weights(rows, n, deadline, clock):
    """Weight histogram (list of n + 1 ints) of the row space of the given independent rows."""
    m = min(len(rows), LOW_ROWS)
    low = [0]
    for g in rows[:m]:
        low += [x ^ g for x in low]
    high = rows[m:]
    hist = [0] * (n + 1)
    acc = Counter()
    h, t, steps = 0, 0, 1 << len(high)
    while True:
        acc.update(map(popcount, map(h.__xor__, low)))
        t += 1
        if t == steps:
            break
        if not t & 15 and clock() > deadline:
            raise Budget
        h ^= high[(t & -t).bit_length() - 1]
    for w, c in acc.items():
        hist[w] += c
    if clock() > deadline:
        raise Budget
    return hist


def correction(G, k):
    s = [((1 if i < k else 0) - popcount(g)) // 2 % 4 for i, g in enumerate(G)]
    cz = [[i, j] for i in range(len(G)) for j in range(i + 1, len(G)) if popcount(G[i] & G[j]) // 2 % 2]
    return s, cz


def hexfrac(fr):
    return {"num": format(fr.numerator, "x"), "den": format(fr.denominator, "x")}


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


def check(I, art, time_budget=None, clock=time.process_time):
    """verify() on an instance given as a loaded dict (tests pass small ad hoc instances this way)."""
    check_instance(I)
    t0 = clock()
    budget = time_budget if time_budget is not None else max(1.0, float(os.environ.get("QARENA_CPU", I["cpu_seconds"])) - 20)
    deadline = t0 + budget
    n, L, S = parse(art, I)
    k, s = len(L), len(S)
    G = L + S
    r = len(G)
    check_rank(G, k)
    check_triorthogonal(G, k)
    lam, sig = check_distance(n, L, S)
    try:
        h0 = enumerate_weights(S, n, deadline, clock)
        hg = enumerate_weights(G, n, deadline, clock)
    except Budget:
        return {"status": "error", "detail": {"reason": "weight enumeration did not finish within the CPU budget "
                                                        "(not ranked)", "n": n, "rank": r}}
    if sum(h0) != 1 << s or sum(hg) != 1 << r or h0[0] != 1 or hg[0] != 1:
        raise RuntimeError("weight enumerator bookkeeping mismatch")
    a, b = I["p_den"] - 2 * I["p_num"], I["p_den"]
    apow, bpow = [1] * (n + 1), [1] * (n + 1)
    for w in range(1, n + 1):
        apow[w] = apow[w - 1] * a
        bpow[w] = bpow[w - 1] * b
    N0 = sum(c * apow[w] * bpow[n - w] for w, c in enumerate(h0) if c)
    N1 = sum(c * apow[w] * bpow[n - w] for w, c in enumerate(hg) if c)
    if N0 <= 0 or N1 < 0 or N1 > N0 << k:
        raise RuntimeError("probability numerators out of range")
    A = Fraction(N0, bpow[n] << s)
    q_bad = Fraction((N0 << k) - N1, N0 << k)
    cost = Fraction(n * bpow[n] << s, k * N0)
    if clock() > deadline:
        return {"status": "error", "detail": {"reason": "verification did not finish within the CPU budget (not ranked)",
                                              "n": n, "rank": r}}
    ceiling = Fraction(I["error_num"], I["error_den"])
    w3 = weight3_logical(n, lam, sig)
    if q_bad > ceiling:
        raise Reject(f"the conditional probability of any output error q_bad = 1 - B/A = {float(q_bad):.6e} exceeds the "
                     f"ceiling {I['error_num']}/{I['error_den']} at p = {I['p_num']}/{I['p_den']} (exact comparison; "
                     f"{w3} accepted weight-3 errors flip an output; A = {float(A):.9f})")
    s_pow, cz = correction(G, k)
    return {"status": "accepted",
            "score": {"cost": hexfrac(cost), "q_bad": hexfrac(q_bad), "n": n, "rank": r},
            "detail": {"n": n, "k": k, "rank": r, "stabilizers": s, "A": hexfrac(A),
                       "A_decimal": f"{float(A):.12f}", "q_bad_decimal": f"{float(q_bad):.11e}",
                       "cost_decimal": f"{float(cost):.10f}", "weight3_logical_errors": w3,
                       "stabilizer_enumerator": h0, "row_space_enumerator": hg,
                       "s": s_pow, "cz": cz,
                       "check": f"full rank {r}, row parities, all pair and triple overlaps even, d_Z >= 3 (all weight-1 "
                                f"and weight-2 errors), exact A and q_bad from the weight enumerators of the "
                                f"{1 << s} stabilizer words and {1 << r} row-space words",
                       "verify_cpu_seconds": round(clock() - t0, 3)}}