# verif_ms_f2.py  --  independent check, plain Python + numpy only.
# Mesirov-Sweet property over F2[T]: for monic g of degree d there is f, gcd(f,g)=1,
# deg f < d, with every partial quotient of f/g of degree 1 or 2.
# k(g) = least number of quotients of degree 2 among such f.
# Part 0: arithmetic self-test.  Part 1: exhaustive recomputation at the largest degree in D_CHOICES
# that fits MEM_BUDGET_MB and BUDGET_S
# (counts of g with k = 0, 1, 2, and none with k >= 3).  Part 2: certificate files
# (if present next to this script) checked by the Euclidean algorithm.
# Fixed parameters, no command line.  Expected run time: well under 10 minutes.
import time, os, operator
import numpy as np

T_START = time.time()
MEM_BUDGET_MB = 40          # memory budget for Part 1 (predicted peak about 64 * 2**D bytes)
D_CHOICES = [22, 20, 18, 16]  # Part 1 uses the largest degree that fits the memory and time budgets
BUDGET_S = 480.0            # time budget for Part 1 (seconds)
RUN_PART1 = True            # switch: set to False to skip the exhaustive part
CERT_FILES = ["certs_d28_two_insertions.csv", "certs_d33_sample.csv"]
EXPECT = {16: (19921, 44251, 1364), 18: (75053, 182509, 4582),
          20: (284755, 748847, 14974), 22: (1085674, 3058561, 50069)}   # counts of k = 0, 1, 2

def say(*a):
    print(*a, flush=True)

XOR = np.bitwise_xor
AND = np.bitwise_and
SHL = np.left_shift
SHR = np.right_shift
U64 = np.uint64

# ---------- Part 0: arithmetic on Python ints (xor only through operator.xor / numpy.bitwise_xor) ----------
def pdeg(a):
    return int(a).bit_length() - 1

def pmul(a, b):
    a = int(a); b = int(b); r = 0; i = 0
    while b:
        if b & 1:
            r = operator.xor(r, a << i)
        b >>= 1; i += 1
    return r

def pdivmod(a, b):
    a = int(a); b = int(b); q = 0; db = pdeg(b)
    while a and pdeg(a) >= db:
        s = pdeg(a) - db
        q = operator.or_(q, 1 << s)
        a = operator.xor(a, b << s)
    return q, a

def cf_from_quotients(qs):
    # f/g = [0; a1,...,ar]  ->  (g, f) via the backward recurrence
    X, Y = 1, 0
    for a in reversed(qs):
        X, Y = operator.xor(pmul(a, X), Y), X
    return X, Y          # g = K(a1..ar), f = K(a2..ar)

def quotients(g, f):
    out = []
    A, B = int(g), int(f)
    while B:
        q, r = pdivmod(A, B)
        out.append(q); A, B = B, r
    return out, A        # A = gcd

def selftest():
    # (T**2+T+1)(T+1) = T**3+1
    assert pmul(0b111, 0b11) == 0b1001
    q, r = pdivmod(0b1001, 0b111)
    assert (q, r) == (0b11, 0)
    qs = [0b10, 0b111, 0b11, 0b101, 0b10]          # T, T**2+T+1, T+1, T**2+1, T
    g, f = cf_from_quotients(qs)
    back, gc = quotients(g, f)
    assert back == qs and gc == 1, "continued fraction round trip failed"
    # vectorised Euclid must accept this pair and reject 1/g (one quotient of degree 7)
    ok = check_pairs(np.array([int(g), int(g)], dtype=U64), np.array([int(f), int(1)], dtype=U64), int(pdeg(g)))[0]
    assert bool(ok[0]) and not bool(ok[1]), "vectorised checker self-test failed"
    say("Part 0: arithmetic self-test passed")

# ---------- vectorised Euclid: all quotient degrees in {1,2}, gcd 1 ----------
def vdeg(x):
    # exact degree for 0 < x < 2**53 via frexp
    m, e = np.frexp(x.astype(np.float64))
    return e.astype(np.int64) - 1

def check_pairs(G, F, d):
    G = G.astype(U64); F = F.astype(U64)
    n = G.shape[0]
    ok = (vdeg(np.maximum(G, U64(1))) == d) & (F > 0)
    ok &= vdeg(np.maximum(F, U64(1))) < d
    nq2 = np.zeros(n, dtype=np.int64)
    A = G.copy(); B = F.copy()
    for _ in range(2 * d + 2):
        act = ok & (B > 0)
        if not act.any():
            break
        da = vdeg(np.maximum(A, U64(1))); db = vdeg(np.maximum(B, U64(1)))
        e = da - db
        bad = act & ((e < 1) | (e > 2))
        ok &= ~bad
        act &= ~bad
        ee = np.where(act, e, 0).astype(U64)
        R = np.where(act, XOR(A, SHL(B, ee)), A)                      # leading term
        dbu = db.astype(U64)
        # next bits of the quotient
        for s in (1, 2):
            sh = (ee - U64(s))
            need = act & (e >= s)
            bitpos = np.where(need, dbu + sh, U64(0))
            hit = need & (AND(SHR(R, bitpos), U64(1)) == U64(1))
            R = np.where(hit, XOR(R, SHL(B, np.where(hit, sh, U64(0)))), R)
        nq2 += (act & (e == 2)).astype(np.int64)
        A = np.where(act, B, A)
        B = np.where(act, R, B)
    ok &= (B == 0) & (A == U64(1))
    return ok, nq2

# ---------- Part 1: exhaustive recomputation at a fixed degree ----------
def cell_images(pattern, d):
    # all g = K(a1..ar) for the cell (list of quotient degrees), as low-d-bit codes
    K1 = np.array([1], dtype=U64); K0 = np.array([0], dtype=U64)
    for e in pattern:
        n = K1.shape[0]
        outs1 = []; outs0 = []
        e = int(e)
        for c in range(int(1) << e):
            a = operator.or_(int(1) << e, int(c))
            prod = np.zeros(n, dtype=U64)
            for b in range(e + int(1)):
                if (a >> b) & int(1):
                    prod = XOR(prod, SHL(K1, U64(int(b))))
            outs1.append(XOR(prod, K0)); outs0.append(K1)
        K1 = np.concatenate(outs1); K0 = np.concatenate(outs0)
    return AND(K1, U64(int((int(1) << int(d)) - int(1))))

def patterns(d, kmax):
    res = []
    def rec(s, cur, k):
        if s == d:
            res.append((k, list(cur))); return
        cur.append(1); rec(s + 1, cur, k); cur.pop()
        if s + 2 <= d and k < kmax:
            cur.append(2); rec(s + 2, cur, k + 1); cur.pop()
    rec(0, [], 0)
    return res

class OverBudget(Exception):
    pass

def part1(d):
    d = int(d)
    best = np.full(int(1) << d, int(255), dtype=np.uint8)
    pats = patterns(d, 2)
    t0 = time.time()
    for i, (k, p) in enumerate(pats):
        img = cell_images(p, d)
        sel = best[img] > k
        best[img[sel]] = k
        if i == 3:
            per = (time.time() - t0) / 4.0
            pred = per * len(pats)
            say("Part 1: %d cells at degree %d, predicted %.0f s" % (len(pats), d, pred))
            if pred > BUDGET_S:
                raise OverBudget()
    cnt = [int((best == k).sum()) for k in (0, 1, 2)]
    none = int((best == 255).sum())
    say("Part 1: degree %d  k=0: %d  k=1: %d  k=2: %d  none with k<=2: %d" % (d, cnt[0], cnt[1], cnt[2], none))
    if d in EXPECT:
        say("Part 1: agreement with the reported counts:", tuple(cnt) == EXPECT[d])
    return none == 0

def run_part1():
    for d in D_CHOICES:
        peak_mb = 64.0 * float(int(1) << int(d)) / 1e6
        if peak_mb > MEM_BUDGET_MB:
            say("Part 1: degree %d skipped (predicted peak %.0f MB > budget %.0f MB)" % (d, peak_mb, MEM_BUDGET_MB))
            continue
        say("Part 1: degree %d, predicted memory peak %.0f MB" % (d, peak_mb))
        try:
            return part1(d)
        except OverBudget:
            say("Part 1: degree %d over the time budget, trying a smaller degree" % d)
    say("Part 1: no degree fits the budgets, skipped")
    return None

def part2():
    here = os.path.dirname(os.path.abspath(__file__))
    for name in CERT_FILES:
        path = os.path.join(here, name)
        if not os.path.exists(path):
            say("Part 2: %s not found, skipped" % name); continue
        arr = np.loadtxt(path, delimiter=",", skiprows=int(1), dtype=np.int64)
        G = arr[:, 0].astype(U64); F = arr[:, 1].astype(U64)
        d = int(vdeg(G[:1])[0])
        ok_all = 0; maxq2 = 0
        for s in range(0, G.shape[0], 1 << 20):
            ok, nq2 = check_pairs(G[s:s + (1 << 20)], F[s:s + (1 << 20)], d)
            ok_all += int(ok.sum()); maxq2 = max(maxq2, int(nq2[ok].max()) if ok.any() else 0)
        say("Part 2: %s  degree %d  certificates %d  valid %d  max number of degree-2 quotients %d"
            % (name, d, G.shape[0], ok_all, maxq2))
        # control: flip one coefficient of f (below degree d) in the first 1000 certificates
        m = min(int(1000), int(G.shape[0]))
        rng = np.random.default_rng(int(20261107))
        flips = SHL(U64(1), rng.integers(0, d, size=m).astype(U64))
        okb, nqb = check_pairs(G[:m], XOR(F[:m], flips), d)
        strict = okb & (nqb <= maxq2)
        say("Part 2: control, %d corrupted certificates: accepted %d (of which with at most %d quotients of degree 2: %d)"
            % (m, int(okb.sum()), maxq2, int(strict.sum())))

selftest()
if RUN_PART1:
    run_part1()
part2()
say("total time %.1f s" % (time.time() - T_START))