# VERIFY_canaday.py
# Ancillary file for the paper "On the largest irreducible factor of an odd perfect
# polynomial over F_2".
# David Niedbala Giraudin, ORCID 0009-0009-1526-1178.
# Technical assistance: Claude (Anthropic).
#
# Plain Python 3 + numpy + python-flint (pip install numpy python-flint).
# Run in plain Python (not under Sage). No external file is read or written.
# Exponents are written with ** throughout; polynomials over F_2 are encoded as
# integers (bit i = coefficient of x**i); addition in F_2[x] is the bitwise xor,
# called through operator.xor / numpy.bitwise_xor.
#
# Usage:
#   python3 VERIFY_canaday.py            quick mode (under a minute):
#                                        sections [0]-[3], Theorem 3 for D = 2..13,
#                                        orbit test for even degrees 4..22
#   python3 VERIFY_canaday.py full       everything in the paper: Theorem 3 for
#                                        D = 2..19 and orbit test for degrees 4..30
#                                        (about two hours on one core)
#   python3 VERIFY_canaday.py general D1 D2     Theorem 3 for D = D1..D2 only
#   python3 VERIFY_canaday.py orbit d1 d2       orbit test for even d in [d1, d2] only
#
# Contents:
#   [0] arithmetic self-tests (vectorised F_2[x] kernels against plain integers and flint)
#   [1] irreducibles of degree <= 19 by a vectorised Rabin test; Gauss's formula
#   [2] brute force over all polynomials of degree <= 16: perfect ones, none odd
#   [3] Lemma 6 (exponent bound) and Lemma 8 (prefilter) tested against direct computation
#   [4] Theorem 3: the set E_D is empty, D = 2..19 (admissible pairs + sieve, Section 3)
#   [5] cross-check of Section 4: no prime of even degree d <= 30 is the factor of smallest
#       degree of a special perfect polynomial (orbit test; Table 2 of the paper)
# Every count printed in sections [4] and [5] is compared with the value printed in the
# paper (tables EXPECTED_GENERAL and EXPECTED_ORBIT below); the script stops on any
# disagreement and ends with the line "ALL CHECKS PASSED".

import operator
import sys
import time

import numpy as np
from flint import nmod_poly

XOR = operator.xor
U64 = np.uint64
ONE = U64(1)
NPX = np.bitwise_xor

# ---------------------------------------------------------------- values of the paper
# D : (|S_D|, admissible pairs, sieve iterations)
EXPECTED_GENERAL = {
    2: (1, 0, 1), 3: (3, 0, 1), 4: (6, 4, 2), 5: (12, 4, 2), 6: (21, 10, 3),
    7: (39, 12, 3), 8: (69, 39, 3), 9: (125, 55, 3), 10: (224, 118, 4),
    11: (410, 164, 4), 12: (745, 418, 4), 13: (1375, 550, 4), 14: (2536, 1086, 5),
    15: (4718, 1682, 5), 16: (8798, 3668, 5), 17: (16508, 5598, 5),
    18: (31040, 12277, 6), 19: (58634, 19351, 6),
}
# d : (primes of degree d congruent to 1 mod x**2+x+1, survivors of the rough check,
#      number of survivors after k orbit levels, k = 1, 2, ...)
EXPECTED_ORBIT = {
    4: (1, 1, [0]), 6: (3, 2, [2, 0]), 8: (10, 4, [4, 0]),
    10: (33, 11, [8, 2, 0]), 12: (111, 24, [20, 6, 0]), 14: (387, 82, [75, 4, 0]),
    16: (1360, 338, [317, 29, 0]), 18: (4842, 894, [843, 62, 0]),
    20: (17459, 2897, [2696, 186, 0]), 22: (63519, 9915, [9335, 526, 0]),
    24: (232950, 33464, [31694, 1721, 12, 0]),
    26: (860265, 110419, [105161, 5155, 8, 0]),
    28: (3195465, 380186, [363023, 16803, 32, 0]),
    30: (11930067, 1322127, [1266736, 53693, 128, 0]),
}
# For d = 28 and d = 30 the first computation of the paper went directly to depth 4; the
# per-level histories above come from the second computation (this script).


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


def check(cond, msg):
    if not cond:
        say("FAILED: " + msg)
        sys.exit(1)


# ---------------------------------------------------------------- plain integer F_2[x]
def deg(a):
    return a.bit_length() - 1


def pmul(a, b):
    r = 0
    while b:
        low = b & -b
        r = XOR(r, a << (low.bit_length() - 1))
        b = XOR(b, low)
    return r


def pdivmod(a, b):
    db = deg(b)
    q = 0
    while deg(a) >= db:
        s = deg(a) - db
        q = q | (1 << s)
        a = XOR(a, b << s)
    return q, a


def pmod(a, b):
    return pdivmod(a, b)[1]


def ppow(a, e):
    r = 1
    while e:
        if e & 1:
            r = pmul(r, a)
        e >>= 1
        if e:
            a = pmul(a, a)
    return r


M1 = 0b111          # x**2 + x + 1
XPOLY = 0b10        # x


def phi3(p):
    return XOR(XOR(pmul(p, p), p), 1)


# ---------------------------------------------------------------- flint bridges
def to_fl(a):
    return nmod_poly([(a >> i) & 1 for i in range(a.bit_length())], 2)


def from_fl(f):
    r = 0
    for i, c in enumerate(f.coeffs()):
        if int(c):
            r = r | (1 << i)
    return r


FX = nmod_poly([0, 1], 2)


def small_primes(n):
    out = []
    q = 2
    while n > 1:
        if n % q == 0:
            out.append(q)
            while n % q == 0:
                n //= q
        q += 1
    return out


def rabin_fl(f):
    """Rabin irreducibility test of a flint polynomial over F_2 (pow_mod and gcd only)."""
    n = f.degree()
    if n < 1:
        return False
    if n == 1:
        return True
    if FX.pow_mod(2 ** n, f) != FX:
        return False
    for q in small_primes(n):
        if (FX.pow_mod(2 ** (n // q), f) - FX).gcd(f).degree() != 0:
            return False
    return True


def factor_certified(f):
    """flint factorisation of f, certified: exact reconstruction + Rabin test of each factor.
    Returns a dict {factor (int): multiplicity}."""
    _, fac = f.factor()
    prod = nmod_poly([1], 2)
    out = {}
    for g, e in fac:
        check(rabin_fl(g), "flint returned a reducible factor")
        prod = prod * g ** e
        out[from_fl(g)] = e
    check(prod == f, "flint factorisation does not reconstruct the polynomial")
    return out


# ---------------------------------------------------------------- vectorised F_2[x], degree <= 62
def v_bitlen(a):
    hi = (a >> U64(32)).astype(np.float64)
    lo = (a & U64(0xFFFFFFFF)).astype(np.float64)
    ehi = np.frexp(hi)[1]
    elo = np.frexp(lo)[1]
    return np.where(hi > 0, 32 + ehi, elo).astype(np.int64)


def v_mulmod(a, b, f, d):
    """a*b mod f for arrays, every f monic of degree d (d <= 62), a, b of degree < d."""
    r = np.zeros_like(f)
    for i in range(d - 1, -1, -1):
        r = r << ONE
        r = NPX(r, f * ((r >> U64(d)) & ONE))
        r = NPX(r, a * ((b >> U64(i)) & ONE))
    return r


def v_gcd(a, b):
    """gcd of the arrays a and b (elementwise), plain Euclid with shifts."""
    a = a.copy()
    b = b.copy()
    for _ in range(140):
        da = v_bitlen(a)
        db = v_bitlen(b)
        sw = da < db
        a, b = np.where(sw, b, a), np.where(sw, a, b)
        da, db = np.where(sw, db, da), np.where(sw, da, db)
        act = b != 0
        if not act.any():
            return a
        sh = np.where(act, da - db, 0).astype(U64)
        a = np.where(act, NPX(a, b << sh), a)
    check(False, "v_gcd did not terminate")


def v_frob(f, d, kmax):
    """list s[k] = x**(2**k) mod f, k = 0..kmax (arrays)."""
    s = [np.full_like(f, XPOLY)]
    for _ in range(kmax):
        s.append(v_mulmod(s[-1], s[-1], f, d))
    return s


def v_irreducible(c, d):
    """Rabin test for an array c of polynomials of degree exactly d (2 <= d <= 62)."""
    s = v_frob(c, d, d)
    ok = s[d] == XPOLY
    prod = np.ones_like(c)
    for q in small_primes(d):
        prod = v_mulmod(prod, NPX(s[d // q], U64(XPOLY)), c, d)
    return ok & (v_gcd(prod, c) == ONE)


def v_square_spread(p, d):
    """p**2 for arrays of degree d (2d <= 62)."""
    r = np.zeros_like(p)
    for i in range(d + 1):
        r = r | (((p >> U64(i)) & ONE) << U64(2 * i))
    return r


def v_residue_mod_m1(c, nbits):
    """(c mod x**2+x+1) as r0 + r1 x, through c(w) = A + B w + C w**2."""
    m = [sum(1 << i for i in range(j, nbits, 3)) for j in range(3)]
    par = [np.bitwise_count(c & U64(mm)).astype(np.int64) & 1 for mm in m]
    a, b, cc = par
    return (a + cc) & 1, (b + cc) & 1


# ---------------------------------------------------------------- [0] self-tests
def section0():
    say("[0] arithmetic self-tests")
    rng = np.random.default_rng(12345)
    import random
    random.seed(12345)
    for d in (5, 17, 30, 44, 60, 62):
        f = rng.integers(0, 2 ** d, size=300, dtype=np.uint64) | U64(1 << d)
        a = rng.integers(0, 2 ** d, size=300, dtype=np.uint64)
        b = rng.integers(0, 2 ** d, size=300, dtype=np.uint64)
        r = v_mulmod(a, b, f, d)
        g = v_gcd(a, f)
        for i in range(300):
            fi, ai, bi = int(f[i]), int(a[i]), int(b[i])
            check(int(r[i]) == pmod(pmul(ai, bi), fi), "v_mulmod")
            check(to_fl(int(g[i])) == to_fl(ai).gcd(to_fl(fi)) or ai == 0, "v_gcd")
            check(to_fl(pmul(ai, bi)) == to_fl(ai) * to_fl(bi), "pmul vs flint")
    c = np.arange(1 << 12, 1 << 13, dtype=np.uint64)
    r0, r1 = v_residue_mod_m1(c, 14)
    for i in range(0, len(c), 7):
        res = pmod(int(c[i]), M1)
        check(res == int(r0[i]) + 2 * int(r1[i]), "residue mod M1")
    say("    v_mulmod, v_gcd, residue mod x**2+x+1, pmul: agree with plain integers and flint")


# ---------------------------------------------------------------- [1] irreducibles
def mobius(n):
    r = 1
    q = 2
    while q * q <= n:
        if n % q == 0:
            n //= q
            if n % q == 0:
                return 0
            r = -r
        q += 1
    return -r if n > 1 else r


def gauss(d):
    return sum(mobius(e) * 2 ** (d // e) for e in range(1, d + 1) if d % e == 0) // d


def irreducibles_of_degree(d):
    if d == 1:
        return [0b10, 0b11]
    c = (np.arange(1 << (d - 1), dtype=np.uint64) << ONE) | U64((1 << d) | 1)
    c = c[(np.bitwise_count(c) & U64(1)) == ONE]
    return sorted(int(v) for v in c[v_irreducible(c, d)])


IRR = {}


def section1(dmax):
    say("[1] irreducible polynomials of degree <= %d (vectorised Rabin test)" % dmax)
    for d in range(1, dmax + 1):
        IRR[d] = irreducibles_of_degree(d)
        check(len(IRR[d]) == gauss(d), "Gauss count at degree %d" % d)
    for d in range(2, 11):
        for p in IRR[d]:
            check(rabin_fl(to_fl(p)), "irreducible list vs flint")
    say("    counts equal to Gauss's formula for d = 1..%d (e.g. N(%d) = %d);"
        " lists cross-checked with flint up to degree 10" % (dmax, dmax, len(IRR[dmax])))


# ---------------------------------------------------------------- sigma via factorisation
def sigma_prime_power(p, h):
    q, r = pdivmod(XOR(ppow(p, h + 1), 1), XOR(p, 1))
    check(r == 0, "sigma division")
    return q


def sigma_int(a):
    fac = factor_certified(to_fl(a))
    s = 1
    for p, e in fac.items():
        s = pmul(s, sigma_prime_power(p, e))
    return s


# ---------------------------------------------------------------- [2] brute force
def section2():
    say("[2] brute force: every non-zero polynomial of degree <= 16")
    found = []
    for a in range(2, 1 << 17):
        if sigma_int(a) == a:
            found.append(a)
    odd = [a for a in found if (a & 1) == 1 and (bin(a).count("1") & 1) == 1]
    say("    perfect polynomials of degree 1..16 found: %d ; odd ones among them: %d"
        % (len(found), len(odd)))
    say("    degrees: %s" % sorted(deg(a) for a in found))
    check(len(found) == 12 and not odd, "brute force")


# ---------------------------------------------------------------- general case helpers
def divisors(n):
    return [e for e in range(1, n + 1) if n % e == 0]


def exponent_candidates(D):
    ts = set()
    for m in range(1, D + 1):
        n = 2 ** m - 1
        e = 1
        while e * e <= n:
            if n % e == 0:
                ts.add(e)
                ts.add(n // e)
            e += 1
    ts.discard(1)
    return sorted(ts)


def euler_phi(n):
    r, m, q = n, n, 2
    while q * q <= m:
        if m % q == 0:
            r -= r // q
            while m % q == 0:
                m //= q
        q += 1
    if m > 1:
        r -= r // m
    return r


def ord2(t):
    v, m = 2 % t, 1
    while v != 1:
        v = (2 * v) % t
        m += 1
    return m


CYC = {}


def cyclo_fl(n):
    """Phi_n(Y) reduced mod 2, as a flint polynomial."""
    if n not in CYC:
        f = FX ** n - 1
        for e in divisors(n):
            if e < n:
                q, r = divmod(f, cyclo_fl(e))
                check(r == 0, "cyclotomic division")
                f = q
        CYC[n] = f
    return CYC[n]


def squarefree_is_smooth(h, D):
    """For h squarefree: h is D-smooth iff h divides prod_{ceil(D/2) <= i <= D} (x**(2**i) - x),
    since every k <= D has a multiple in [ceil(D/2), D]."""
    s = FX % h
    prod = nmod_poly([1], 2) % h
    for i in range(1, D + 1):
        s = (s * s) % h
        if 2 * i >= D:
            prod = (prod * (s - FX)) % h
    return prod == 0


def smooth_factor(g, D, irrsets):
    """Returns (dict factor -> multiplicity, True) if g is D-smooth, (None, False) otherwise.
    Non-smoothness is certified by a squarefree divisor h of g (h | g and gcd(h, h') = 1
    checked) failing the test above; smoothness by a flint factorisation whose product is
    checked to equal g and whose factors all belong to the list of irreducibles of [1]."""
    for h, _ in g.factor_squarefree()[1]:
        if h.degree() == 0:
            continue
        if not squarefree_is_smooth(h, D):
            check(g % h == 0, "squarefree part does not divide g")
            check(h.gcd(h.derivative()).degree() == 0, "squarefree part is not squarefree")
            return None, False
    _, fac = g.factor()
    out = {}
    prod = nmod_poly([1], 2)
    for u, e in fac:
        ui = from_fl(u)
        check(u.degree() <= D and ui in irrsets[u.degree()], "smooth factor not in the list")
        out[ui] = e
        prod = prod * u ** e
    check(prod == g, "factorisation does not reconstruct g")
    return out, True


def admissible_pairs(D, use_l4=True):
    S = [p for d in range(2, D + 1) for p in IRR[d]]
    irrsets = {d: set(IRR[d]) for d in range(1, D + 1)}
    ts = exponent_candidates(D)
    tset = set(ts)
    for t in ts:
        for e in divisors(t):
            check(e == 1 or e in tset, "T_D closed under divisors")
    N2 = {e: len(IRR[e]) for e in range(1, D + 1)}
    l4 = {}
    for t in ts:
        m = ord2(t)
        l4[t] = (m, euler_phi(t), sum(j * N2[j] for j in range(m, D + 1, m)))
    pairs = {}
    n_eval = 0
    n_cut = 0
    for P in S:
        p = deg(P)
        fP = to_fl(P)
        smooth = {}
        alive = set(ts)
        for t in ts:
            if not any(u % t == 0 for u in alive):
                continue
            m, ph, B = l4[t]
            if use_l4:
                extra = p - 1 if m <= p - 1 else 0
                if ph * p > B + extra:
                    n_cut += 1
                    alive = {u for u in alive if u % t != 0}
                    continue
            g = cyclo_fl(t)(fP)
            n_eval += 1
            fac, ok = smooth_factor(g, D, irrsets)
            if not ok:
                alive = {u for u in alive if u % t != 0}
            else:
                check(0b10 not in fac and 0b11 not in fac, "Phi_t(P) divisible by x or x+1")
                smooth[t] = fac
        for t in sorted(alive):
            prod = {}
            for e in divisors(t):
                if e > 1:
                    for R, mult in smooth[e].items():
                        prod[R] = prod.get(R, 0) + mult
            pairs[(P, t - 1)] = prod
    return S, pairs, n_eval, n_cut


def sieve(S, pairs):
    """Rules (R1)-(R3) of Section 3, iterated to a fixed point. Returns (alive primes, iterations)."""
    alive_p = set(S)
    alive_pairs = dict(pairs)
    it = 0
    while True:
        it += 1
        before = len(alive_pairs)
        alive_pairs = {k: v for k, v in alive_pairs.items()
                       if k[0] in alive_p and all(R in alive_p for R in v)}
        killed_pairs = before - len(alive_pairs)
        has_pair = {k[0] for k in alive_pairs}
        produced = {R for v in alive_pairs.values() for R, e in v.items() if e > 0}
        new = {P for P in alive_p if P in has_pair and P in produced}
        if new == alive_p and killed_pairs == 0:
            break
        alive_p = new
        if not alive_p:
            break
    return alive_p, it


def section3():
    say("[3] Lemma 6 and Lemma 8 against direct computation")
    for D in range(2, 7):
        S = [p for d in range(2, D + 1) for p in IRR[d]]
        irrsets = {d: set(IRR[d]) for d in range(1, D + 1)}
        tset = set(exponent_candidates(D))
        direct = {}
        for P in S:
            for h in range(2, 2 ** D - 1, 2):
                g = to_fl(sigma_prime_power(P, h))
                fac, ok = smooth_factor(g, D, irrsets)
                if ok and 0b10 not in fac and 0b11 not in fac:
                    check(h + 1 in tset, "Lemma 6 violated")
                    direct[(P, h)] = fac
        _, pairs, _, _ = admissible_pairs(D)
        check(direct == pairs, "cyclotomic route vs direct sigma, D = %d" % D)
    say("    D <= 6: every even h <= 2**D - 2 tried by direct sigma; admissible pairs and"
        " production vectors identical to the cyclotomic route; no h outside T_D")
    for D in range(2, 13):
        _, a, _, _ = admissible_pairs(D, use_l4=True)
        _, b, _, _ = admissible_pairs(D, use_l4=False)
        check(a == b, "Lemma 8 prefilter changes the admissible pairs at D = %d" % D)
    say("    D <= 12: admissible pairs with and without the prefilter of Lemma 8: identical")


def section4(D1, D2):
    say("[4] Theorem 3: E_D is empty")
    for D in range(D1, D2 + 1):
        t0 = time.time()
        S, pairs, n_eval, n_cut = admissible_pairs(D)
        for (P, h), v in pairs.items():
            check(P not in v, "a pair produces its own prime")
        alive, it = sieve(S, pairs)
        say("    D = %2d : |S_D| = %6d, admissible pairs = %6d, sieve iterations = %d,"
            " survivors = %d   (%d evaluations, %d prefilter cuts, %.0f s)"
            % (D, len(S), len(pairs), it, len(alive), n_eval, n_cut, time.time() - t0))
        check(not alive, "the sieve does not empty at D = %d" % D)
        check((len(S), len(pairs), it) == EXPECTED_GENERAL[D], "Table 1 of the paper, D = %d" % D)


# ---------------------------------------------------------------- [5] orbit test
def sweep(d, batch_bits=21):
    """Primes P of degree d with P = 1 mod x**2+x+1, and those for which Phi_3(P) has no
    irreducible factor of degree < d (rough check). Fully vectorised, no flint."""
    n_cong = 0
    kept = []
    total = 1 << (d - 1)
    step = min(total, 1 << batch_bits)
    top = U64((1 << d) | 1)
    evens = list(range(2, d - 1, 2))
    for lo in range(0, total, step):
        c = (np.arange(lo, lo + step, dtype=np.uint64) << ONE) | top
        r0, r1 = v_residue_mod_m1(c, d + 1)
        c = c[(r0 == 1) & (r1 == 0)]
        c = c[(np.bitwise_count(c) & U64(1)) == ONE]
        if len(c) == 0:
            continue
        c = c[v_irreducible(c, d)]
        n_cong += len(c)
        if len(c) == 0:
            continue
        f = NPX(NPX(v_square_spread(c, d), c), ONE)
        if evens:
            s = v_frob(f, 2 * d, evens[-1])
            prod = np.ones_like(f)
            for e in evens:
                prod = v_mulmod(prod, NPX(s[e], U64(XPOLY)), f, 2 * d)
            c = c[v_gcd(prod, f) == ONE]
        kept.extend(int(v) for v in c)
    return n_cong, kept


def deep(d, survivors, depth_cap=12):
    memo = {}

    def level(R):
        fac = factor_certified(to_fl(phi3(R)))
        kids = []
        for Q in fac:
            if deg(Q) % 2 or deg(Q) < d or pmod(Q, M1) != 1:
                return None
            kids.append(Q)
        return kids

    lev = {}

    def ok(R, k):
        if k == 0:
            return True
        key = (R, k)
        if key not in memo:
            if R not in lev:
                lev[R] = level(R)
            kids = lev[R]
            memo[key] = kids is not None and all(ok(Q, k - 1) for Q in kids)
        return memo[key]

    alive = list(survivors)
    hist = []
    k = 0
    while alive and k < depth_cap:
        k += 1
        alive = [P for P in alive if ok(P, k)]
        hist.append(len(alive))
    return hist, alive


def selftest_orbit():
    """the vectorised sweep against a direct flint computation, degrees 4..16."""
    for d in range(4, 17, 2):
        n_cong, kept = sweep(d)
        ref_c, ref_k = 0, []
        for P in IRR[d] if d in IRR else irreducibles_of_degree(d):
            if pmod(P, M1) != 1:
                continue
            ref_c += 1
            fac = factor_certified(to_fl(phi3(P)))
            if all(deg(Q) >= d for Q in fac):
                ref_k.append(P)
        check(n_cong == ref_c and sorted(kept) == sorted(ref_k), "sweep vs flint, d = %d" % d)
        for P in ref_k[:50]:
            for Q in factor_certified(to_fl(phi3(P))):
                check(deg(Q) % 2 == 0, "odd-degree factor of Phi_3(P) (parity of degrees, Section 4)")
    say("    vectorised sweep identical to a direct flint computation for d = 4..16")


def section5(d1, d2):
    say("[5] Section 4: orbit test (smallest prime of a special perfect polynomial)")
    selftest_orbit()
    for d in range(d1, d2 + 1, 2):
        t0 = time.time()
        n_cong, kept = sweep(d)
        t1 = time.time()
        hist, alive = deep(d, kept)
        say("    d = %2d : primes = 1 mod x**2+x+1 : %8d ; rough survivors : %7d ;"
            " survivors by orbit depth : %s   (%.0f s + %.0f s)"
            % (d, n_cong, len(kept), hist, t1 - t0, time.time() - t1))
        check(not alive, "orbit not extinguished at d = %d" % d)
        exp = EXPECTED_ORBIT[d]
        check((n_cong, len(kept)) == exp[:2], "orbit table, d = %d" % d)
        check(exp[2] is None or hist == exp[2], "orbit depth history, d = %d" % d)
        check(exp[2] is not None or len(hist) <= 4, "orbit depth > 4 at d = %d" % d)


# ---------------------------------------------------------------- main
def main():
    args = sys.argv[1:]
    mode = args[0] if args else "quick"
    t0 = time.time()
    section0()
    if mode == "general":
        D1, D2 = int(args[1]), int(args[2])
        section1(max(D2, 6))
        section4(D1, D2)
    elif mode == "orbit":
        section1(16)
        section5(int(args[1]), int(args[2]))
    else:
        full = mode == "full"
        section1(19 if full else 16)
        section2()
        section3()
        section4(2, 19 if full else 13)
        section5(4, 30 if full else 22)
    say("ALL CHECKS PASSED (%.0f s)" % (time.time() - t0))


if __name__ == "__main__":
    main()