"""Board qldpc: weight-6 CSS codes, maximize k at a rung (n <= n_max, certified d >= d_min). Catalog W1-4 (Q18).

Artifact {"n": n, "HX": [[q, ...], ...], "HZ": [[q, ...], ...]}: rows are strictly increasing qubit lists.
Validity: row weight <= 6; qubit degree (rows of HX and HZ together containing it) <= 6; HX HZ^T = 0 (mod 2);
k = n - rank HX - rank HZ >= 1; n <= n_max; and min(d_X, d_Z) >= d_min CERTIFIED by exhaustive search:

Distance certificate (exact, no solver). A minimum-weight nontrivial X logical x* contains no smaller nonzero
element of ker HZ: such a z would be nontrivial (smaller logical) or a stabilizer (then x* + z is a smaller
nontrivial logical). So a depth-first search finds x* if it has weight <= w: start from each qubit q0 (forbidding
lower qubits), repeatedly take the lowest unsatisfied Z check, and branch on which of its qubits outside S and the
forbidden set to add (branch i forbids the qubits of branches 1..i-1). Stop a branch when the syndrome is zero
(then test nontriviality against a basis of Z logicals), when |S| = w, or when the unsatisfied checks cannot all be
fixed by w - |S| more qubits (cover bound below). Along the branch that follows x*, S stays a proper subset of x*, so
the syndrome is nonzero and the chosen check meets x* outside S (even overlap with x*, odd with S). If no branch finds
a nontrivial zero-syndrome set with w = d_min - 1, then d_X >= d_min. Same for Z with the roles of HX and HZ swapped.

Cover bound. On the branch that follows x*, the rest T = x* - S avoids S and the forbidden set, |T| <= w - |S|, and T
meets every unsatisfied check c (odd overlap). Let a_q be the number of unsatisfied checks containing a free qubit q
and m_c = max of a_q over the free qubits of c (m_c >= a_q for every free q in c). Then
    sum_c 1/m_c <= sum_c |c & T| / m_c = sum_{q in T} sum_{unsatisfied c containing q} 1/m_c <= |T|,
because each inner sum has a_q terms, each at most 1/a_q. So a branch with sum_c 1/m_c > w - |S|, or with an
unsatisfied check that has no free qubit, holds no such x* and is cut (the test runs when more checks are unsatisfied
than qubits remain; otherwise the branch dies at that check anyway). The cut never removes a zero-syndrome set (it
removes only subtrees that hold none), so the search visits the same zero-syndrome sets in the same order as without
it and returns the same witness.
Score: k (max), then n (min), then max qubit degree (min).
"""
import time
from math import gcd
from common import Reject, exact_keys, want_int, want_list, is_int, load_instance, rank, rref, reduce, nullspace, popcount, short


def parse_rows(rows, n, name, max_w):
    out, seen = [], set()
    for i, r in enumerate(rows):
        if not isinstance(r, list) or not r:
            raise Reject(f"{name}[{i}] must be a nonempty list of qubit indices")
        if len(r) > max_w:
            raise Reject(f"{name}[{i}] has weight {len(r)} > {max_w}")
        prev = -1
        for q in r:
            if not is_int(q) or not 0 <= q < n:
                raise Reject(f"{name}[{i}] contains {short(q)}, not a qubit index in [0, {n})")
            if q <= prev:
                raise Reject(f"{name}[{i}] must be strictly increasing")
            prev = q
        b = sum(1 << q for q in r)
        if b in seen:
            raise Reject(f"{name}[{i}] repeats an earlier row")
        seen.add(b)
        out.append(b)
    return out


def logical_basis(stab_same, stab_other, n):
    """Basis of ker(stab_other) modulo rowspace(stab_same): the logicals of the type of stab_same."""
    ker = nullspace(stab_other, n)
    basis = dict(rref(stab_same))
    out = []
    for v in ker:
        w = reduce(v, basis)
        if w:
            out.append(v)
            p = w.bit_length() - 1
            for q in list(basis):
                if basis[q] >> p & 1:
                    basis[q] ^= w
            basis[p] = w
    return out


def cover_exceeds(checks, colw, L, syn, used, r):
    """True if the unsatisfied checks syn cannot be fixed by r qubits outside used (cover bound in the docstring).
    colw: largest number of checks on one qubit; L: a common multiple of 1..colw."""
    free = ~used
    lv = [0] * (colw + 1)                                                  # lv[t]: free qubits in >= t unsatisfied checks
    cs = []
    s = syn
    while s:
        low = s & -s
        s ^= low
        a = checks[low.bit_length() - 1] & free
        if not a:
            return True
        cs.append(a)
        for t in range(colw, 1, -1):
            lv[t] |= lv[t - 1] & a
        lv[1] |= a
    tot, cap = 0, L * r
    for a in cs:
        t = colw
        while t > 1 and not lv[t] & a:                                     # t = m_c
            t -= 1
        tot += L // t
        if tot > cap:
            return True
    return False


def search(checks, logicals, n, w, deadline):
    """Return a nontrivial vector x (bitset) with checks.x = 0 and |x| <= w, or None. Exhaustive.
    deadline: time.process_time() value after which TimeoutError is raised.
    checks: list of row bitsets (the stabilizers of the OTHER type); logicals: basis of the conjugate logicals."""
    nc = len(checks)
    chk_q = [[q for q in range(n) if c >> q & 1] for c in checks]          # qubits of each check, ascending
    flip = [0] * n                                                         # checks containing qubit q
    for ci, c in enumerate(checks):
        for q in chk_q[ci]:
            flip[q] |= 1 << ci
    sig = [0] * n                                                          # logical signature of qubit q
    for j, l in enumerate(logicals):
        for q in range(n):
            if l >> q & 1:
                sig[q] |= 1 << j
    colw = max((popcount(f) for f in flip), default=0)
    pc = getattr(int, "bit_count", popcount)                              # int.bit_count is Python >= 3.10
    L = 1                                                                  # common denominator of the terms 1/m_c
    for t in range(2, colw + 1):
        L = L * t // gcd(L, t)
    nodes = [0]
    found = [None]

    def dfs(S, size, syn, lsig, forb):
        nodes[0] += 1
        if nodes[0] & 0xFFFFF == 0 and time.process_time() > deadline:   # polled on every node count, pruned or not
            raise TimeoutError
        if syn == 0:
            if lsig:
                found[0] = S
                return True
            return False
        r = w - size
        sw = pc(syn)
        if r == 0 or sw > colw * r:
            return False
        used = S | forb
        if sw > r and cover_exceeds(checks, colw, L, syn, used, r):      # sw <= r: the bound cannot exceed r
            return False
        low = syn & -syn
        ci = low.bit_length() - 1
        f = forb
        for q in chk_q[ci]:
            bit = 1 << q
            if used & bit:
                continue
            if dfs(S | bit, size + 1, syn ^ flip[q], lsig ^ sig[q], f):
                return True
            f |= bit
            used |= bit
        return False

    for q0 in range(n):
        if dfs(1 << q0, 1, flip[q0], sig[q0], (1 << q0) - 1):
            return found[0], nodes[0]
    return None, nodes[0]


def verify(inst, art, time_budget=None):
    I = load_instance("qldpc", inst)
    n_max, d_min, max_w, max_deg = I["n_max"], I["d_min"], I["max_weight"], I["max_degree"]
    exact_keys(art, ["n", "HX", "HZ"])
    n = want_int(art["n"], 1, n_max, f"n (rung {inst} allows n <= {n_max})")
    HXl = want_list(art["HX"], "HX", max_len=2 * n)
    HZl = want_list(art["HZ"], "HZ", max_len=2 * n)
    if len(HXl) + len(HZl) > 2 * n:
        raise Reject("more than 2n rows in total")
    HX = parse_rows(HXl, n, "HX", max_w)
    HZ = parse_rows(HZl, n, "HZ", max_w)
    deg = [0] * n
    for r in HX + HZ:
        for q in range(n):
            if r >> q & 1:
                deg[q] += 1
    maxdeg = max(deg)
    if maxdeg > max_deg:
        q = deg.index(maxdeg)
        raise Reject(f"qubit {q} lies in {maxdeg} checks > {max_deg}")
    for i, a in enumerate(HX):
        for j, b in enumerate(HZ):
            if popcount(a & b) & 1:
                raise Reject(f"HX[{i}] and HZ[{j}] anticommute")
    rx, rz = rank(HX), rank(HZ)
    k = n - rx - rz
    if k < 1:
        raise Reject(f"k = {k}: no logical qubits")
    zbar = logical_basis(HZ, HX, n)   # Z logicals: ker HX mod rowspace HZ
    xbar = logical_basis(HX, HZ, n)   # X logicals: ker HZ mod rowspace HX
    assert len(zbar) == k and len(xbar) == k
    w = d_min - 1
    import os
    budget = time_budget or float(os.environ.get("QARENA_CPU", "21600")) - 60
    deadline = time.process_time() + budget                 # the board's budget is CPU seconds
    t0 = time.time()
    try:
        # cheap shallow pass first: low-weight witnesses (e.g. an unprotected qubit) are found at once
        for ws in range(1, min(3, w + 1)):
            for chk, lg, kind, st in ((HZ, zbar, "X", HX), (HX, xbar, "Z", HZ)):
                x, _ = search(chk, lg, n, ws, deadline)
                if x is not None:
                    return reject_witness(kind, x, n, chk, st)
        x, nodes_x = search(HZ, zbar, n, w, deadline)      # X logicals commute with Z checks
        if x is not None:
            return reject_witness("X", x, n, HZ, HX)
        z, nodes_z = search(HX, xbar, n, w, deadline)
        if z is not None:
            return reject_witness("Z", z, n, HX, HZ)
    except TimeoutError:
        return {"status": "error", "detail": {"reason": "distance search did not finish within the time limit (not ranked)",
                                              "n": n, "k": k}}
    return {"status": "accepted",
            "score": {"k": k, "n": n, "maxdeg": maxdeg, "d_cert": d_min, "kd2_over_n": round(k * d_min * d_min / n, 4)},
            "detail": {"n": n, "k": k, "rank_HX": rx, "rank_HZ": rz, "max_qubit_degree": maxdeg,
                       "max_check_weight": max(popcount(r) for r in HX + HZ),
                       "distance": f"certified d >= {d_min}: exhaustive search found no nontrivial X or Z logical of weight <= {w}",
                       "search_nodes": {"X": nodes_x, "Z": nodes_z}, "search_seconds": round(time.time() - t0, 2)}}


def reject_witness(kind, x, n, checks, stabs):
    # independent recheck of the witness before using it
    assert all(popcount(c & x) % 2 == 0 for c in checks)
    assert reduce(x, rref(stabs)) != 0
    qs = [q for q in range(n) if x >> q & 1]
    raise Reject(f"distance below the rung: nontrivial {kind} logical of weight {len(qs)} on qubits {qs}")