# VERIFY_mersenne7.py
# Ancillary file for the paper "The invariant c_M of Gallardo and Rahavandrainy at Mersenne
# prime exponents, and the case p = 7".
# David Niedbala Giraudin, ORCID 0009-0009-1526-1178.
# Technical assistance: Claude (Anthropic).
#
# Plain Python 3 + python-flint (pip install python-flint). Run in plain Python (not under
# Sage). No external file is read or written. Exponents are written with **.
#
# Usage:
#   python3 VERIFY_mersenne7.py          quick mode (under a minute): sections [0]-[1],
#                                        Table 1 rows p <= 13, section [4] up to degree 120
#   python3 VERIFY_mersenne7.py full     everything in the paper: all rows of Table 1 and
#                                        section [4] up to degree 500 (less than half an hour)
#
# Notation (Gallardo-Rahavandrainy; their Sigma_p with superscript j is written Sigma_p(j)): M is a Mersenne prime, i.e. an irreducible polynomial
# x**a (x+1)**b + 1 with a, b >= 1; p is an odd prime, 2h = p - 1, S = sigma(M**(p-1)) =
# Phi_p(M), U = sigma(S), W = U + S + 1, c_M = deg S - deg W, m_M = least odd l with
# alpha_l(W) = 1 (when deg W is even), alpha_l(F) = coefficient of x**(deg F - l) in F.
#
# Contents:
#   [0] self-tests of the irreducibility and Mersenne tests
#   [1] Theorem 2: c_M = d + 1 if deg M != d and c_M = d if deg M = d, for p = 2**d - 1
#       (p = 7, 31, 127, 8191), with the two facts used in the proof (all irreducibles of
#       degree d other than M divide S; every irreducible factor of S has degree divisible
#       by d), Lemma 6 (parity of the number of factors of smallest degree) and
#       Corollary 3 (Sigma_p(1) = Mersenne primes of degree d or d+1, Sigma_p(2) empty)
#   [3] Table 1 of Gallardo-Rahavandrainy reproduced (counts of Sigma_p(j), 5 <= deg M <= d)
#   [4] p = 7: every Mersenne prime M of degree <= N: c_M, m_M, membership in Sigma_7,
#       and a non-Mersenne irreducible factor of S found directly
# The script stops on any disagreement and ends with the line "ALL CHECKS PASSED".

import sys
import time

from flint import nmod_poly

X = nmod_poly([0, 1], 2)
X1 = nmod_poly([1, 1], 2)
ONE = nmod_poly([1], 2)
ZERO = nmod_poly([0], 2)

# Table 1 of [GR25, Mat. Vesnik]: p : (degree bound, #Sigma_p(1..4) with 5 <= deg M <= bound,
# number of Mersenne primes with 5 <= deg M <= bound)
GR_TABLE = {
    5: (100, [4, 0, 0, 0], 226), 7: (100, [0, 0, 2, 2], 226), 11: (100, [2, 0, 2, 0], 226),
    13: (100, [0, 2, 0, 0], 226), 17: (100, [2, 0, 0, 0], 226), 19: (100, [0, 0, 0, 0], 226),
    23: (100, [0, 0, 0, 0], 226), 29: (100, [4, 0, 2, 0], 226), 53: (60, [4, 0, 4, 0], 138),
    59: (60, [0, 0, 0, 0], 138), 61: (60, [4, 2, 0, 0], 138), 67: (60, [0, 0, 0, 0], 138),
    71: (60, [0, 2, 0, 0], 138),
}
# Section [4]: degree bound : (number of Mersenne primes of degree <= bound, maximum of m_M,
# members of Sigma_7 as (deg M, a, j) with M = x**a (x+1)**(deg M - a) + 1 in Sigma_7(j))
P7_MEMBERS = [(3, 1, 1), (3, 2, 1), (4, 1, 1), (4, 3, 1),
              (6, 1, 4), (6, 5, 4), (7, 3, 3), (7, 4, 3)]
P7_EXPECTED = {120: (281, 11), 500: (1203, 17)}


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


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


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 is_irreducible(f):
    """Rabin test over F_2 (pow_mod and gcd only)."""
    n = f.degree()
    if n < 1:
        return False
    if n == 1:
        return True
    if X.pow_mod(2 ** n, f) != X:
        return False
    return all((X.pow_mod(2 ** (n // q), f) - X).gcd(f).degree() == 0 for q in small_primes(n))


def is_mersenne_prime(P):
    """P irreducible of the form x**a (x+1)**b + 1 with a, b >= 1."""
    Q = P + 1
    a = b = 0
    while Q.degree() > 0 and int(Q.coeffs()[0]) == 0:
        Q = divmod(Q, X)[0]
        a += 1
    while Q.degree() > 0:
        q, r = divmod(Q, X1)
        if r != 0:
            break
        Q = q
        b += 1
    return Q == ONE and a >= 1 and b >= 1


def mersenne_primes(n):
    """the Mersenne primes of degree n, as (a, M) with M = x**a (x+1)**(n-a) + 1."""
    out = []
    for a in range(1, n):
        M = X ** a * X1 ** (n - a) + 1
        if is_irreducible(M):
            out.append((a, M))
    return out


def factor_certified(f):
    _, fac = f.factor()
    prod = ONE
    for g, e in fac:
        check(is_irreducible(g), "flint returned a reducible factor")
        prod = prod * g ** e
    check(prod == f, "flint factorisation does not reconstruct the polynomial")
    return fac


def alpha(F, l):
    c = F.coeffs()
    dg = F.degree()
    return int(c[dg - l]) if 0 <= l <= dg else 0


def invariants(M, p, certify=True):
    """S, its factorisation, c_M and m_M (m_M = None if deg W is odd)."""
    S = ZERO
    for i in range(p):
        S = S + M ** i
    fac = factor_certified(S) if certify else S.factor()[1]
    check(all(e == 1 for _, e in fac), "S is not squarefree")
    U = ONE
    for g, _ in fac:
        U = U * (g + 1)
    W = U + S + 1
    c = S.degree() - W.degree()
    m = None
    if W.degree() % 2 == 0:
        m = next((l for l in range(1, W.degree() + 1, 2) if alpha(W, l)), None)
    return S, [g for g, _ in fac], c, m


def sigma_classes(n, c, m):
    """membership of M (deg M = n) in Sigma_p(1), ..., Sigma_p(4)."""
    return [c == n,
            c == n + 1 and c % 2 == 1,
            m is not None and c + m == n and c % 2 == 0,
            m is not None and c + m == n + 1 and c % 2 == 0 and m >= 3]


# ---------------------------------------------------------------- [0]
def section0():
    say("[0] self-tests")
    import random
    random.seed(7)
    for _ in range(400):
        n = random.randint(2, 40)
        f = nmod_poly([random.randint(0, 1) for _ in range(n)] + [1], 2)
        fac = f.factor()[1]
        ref = len(fac) == 1 and fac[0][1] == 1
        check(is_irreducible(f) == ref, "Rabin test against flint")
    counts = {}
    for n in range(2, 13):
        counts[n] = len(mersenne_primes(n))
        for a, M in mersenne_primes(n):
            check(is_mersenne_prime(M), "Mersenne test")
    say("    Rabin test agrees with flint on 400 random polynomials; Mersenne primes of degree"
        " 2..12: %s" % [counts[n] for n in range(2, 13)])


# ---------------------------------------------------------------- [1]
def irreducibles_of_degree(d):
    out = []
    for k in range(2 ** d):
        f = nmod_poly([(k >> i) & 1 for i in range(d)] + [1], 2)
        if is_irreducible(f):
            out.append(f)
    return out


def section1():
    say("[1] Theorem 2 and the facts used in its proof")
    cases = [(3, 7, 40), (5, 31, 24), (7, 127, 12), (13, 8191, 3)]
    for d, p, nmax in cases:
        irr_d = irreducibles_of_degree(d)
        check(len(irr_d) == (p - 1) // d, "number of irreducibles of degree d")
        tested = 0
        classes = {}
        for n in range(2, nmax + 1):
            for a, M in mersenne_primes(n):
                S, fac, c, m = invariants(M, p, certify=(p < 8191))
                check(all(g.degree() % d == 0 for g in fac), "factor degree not divisible by d")
                small = [g for g in fac if g.degree() == d]
                expected_small = [g for g in irr_d if g != M]
                check(sorted(str(g) for g in small) == sorted(str(g) for g in expected_small),
                      "degree-d factors of S")
                kappa = len(small)
                lemma2_c = d if kappa % 2 == 1 else None
                check(c == (d if n == d else d + 1), "Theorem 2, p = %d, deg M = %d" % (p, n))
                check(lemma2_c is None or c == lemma2_c, "Lemma 6")
                check((kappa % 2 == 1) == (n == d), "parity of kappa")
                cl = [j + 1 for j, v in enumerate(sigma_classes(n, c, m)) if v]
                for j in cl:
                    classes.setdefault(j, []).append(n)
                if 1 in cl:
                    check(n in (d, d + 1), "Corollary 3: Sigma_p(1)")
                check(2 not in cl, "Corollary 3: Sigma_p(2) is empty")
                tested += 1
        say("    p = %4d (d = %2d): %3d Mersenne primes M with 2 <= deg M <= %d, all with"
            " c_M = %d if deg M = %d and c_M = %d otherwise; degrees of the members of"
            " Sigma_p(j): %s" % (p, d, tested, nmax, d, d, d + 1,
                                 {j: sorted(set(v)) for j, v in sorted(classes.items())}))


# ---------------------------------------------------------------- [3]
def section3(primes):
    say("[3] Table 1 of Gallardo-Rahavandrainy")
    cache = {}
    for p in primes:
        bound, exp_counts, exp_total = GR_TABLE[p]
        t0 = time.time()
        counts = [0, 0, 0, 0]
        total = 0
        for n in range(5, bound + 1):
            if n not in cache:
                cache[n] = mersenne_primes(n)
            for a, M in cache[n]:
                total += 1
                _, _, c, m = invariants(M, p, certify=False)
                for j, v in enumerate(sigma_classes(n, c, m)):
                    if v:
                        counts[j] += 1
        say("    p = %2d, 5 <= deg M <= %3d : #Sigma_p(1..4) = %s, Mersenne primes = %d   (%.0f s)"
            % (p, bound, counts, total, time.time() - t0))
        check(counts == exp_counts and total == exp_total, "Table 1, p = %d" % p)


# ---------------------------------------------------------------- [4]
def section4(N):
    say("[4] p = 7, every Mersenne prime of degree <= %d" % N)
    t0 = time.time()
    total = 0
    mmax = 0
    first = {}
    members = []
    for n in range(2, N + 1):
        for a, M in mersenne_primes(n):
            total += 1
            S, fac, c, m = invariants(M, 7, certify=True)
            check(c == (3 if n == 3 else 4), "Theorem 2 for p = 7")
            check(any(not is_mersenne_prime(g) for g in fac),
                  "no non-Mersenne factor for deg M = %d, a = %d" % (n, a))
            if m is not None:
                mmax = max(mmax, m)
                first.setdefault(m, n)
            for j, v in enumerate(sigma_classes(n, c, m)):
                if v:
                    members.append((n, a, j + 1))
    say("    %d Mersenne primes; each S = sigma(M**6) has a non-Mersenne irreducible factor"
        % total)
    say("    largest m_M = %d ; smallest degree at which m_M takes each value: %s"
        % (mmax, dict(sorted(first.items()))))
    say("    members of Sigma_7 (deg M, a, j): %s   (%.0f s)" % (members, time.time() - t0))
    check(members == P7_MEMBERS, "members of Sigma_7")
    exp = P7_EXPECTED.get(N)
    if exp and exp[0] is not None:
        check((total, mmax) == exp, "counts of section [4]")


def main():
    full = len(sys.argv) > 1 and sys.argv[1] == "full"
    t0 = time.time()
    section0()
    section1()
    section3(sorted(GR_TABLE) if full else [5, 7, 11, 13])
    section4(500 if full else 120)
    say("ALL CHECKS PASSED (%.0f s)" % (time.time() - t0))


if __name__ == "__main__":
    main()