#!/usr/bin/env python3
"""
A400659 -- self-contained verification program, Python 3.8 or later.

    a(n) = least k >= 1 such that A002110(n-1)*prime(n)^k - 1 is prime.

Run:    python A400659.py         (checks n = 1..100, the b-file range)
        python A400659.py 30      (n = 1..30 only)

Standard library only. If the gmpy2 package is installed it is used, which is
about ten times faster on the largest terms; the results are the same.

For every n the program finds a(n) itself and proves it:

 * every k < a(n): N = A002110(n-1)*prime(n)^k - 1 is shown composite, either
   by a prime factor p < 10^5 (with p < N), or by a Fermat witness b, i.e.
   b^(N-1) != 1 (mod N), which no prime N can satisfy.
   (N = 1, at n = 1, k = 1, is not prime by definition.)

 * k = a(n): N is proved prime by the Lucas N+1 test, which applies because
   N + 1 = A002110(n-1)*prime(n)^k is completely factored.

The Lucas N+1 test used (Lucas; Lehmer; Brillhart-Lehmer-Selfridge 1975).
Let N >= 3 be odd and let F = N + 1 be completely factored. Suppose there is
one integer D with gcd(D, N) = 1 and Jacobi(D, N) = -1, and for each prime
q dividing F an integer P_q, P_q == D (mod 2), Q_q = (P_q^2 - D)/4, with
    gcd(2*Q_q, N) = 1,
    U_{N+1}(P_q, Q_q) == 0 (mod N),
    gcd(U_{(N+1)/q}(P_q, Q_q), N) = 1,
where U is the Lucas sequence U_0 = 0, U_1 = 1, U_{m+1} = P*U_m - Q*U_{m-1}.
Then every prime r | N satisfies (N+1) | r - Jacobi(D, r), so r >= N and N is
prime. The discriminant D must be the SAME for every q (only P may change);
that is what makes the conditions for the different q combine.

U_m is read off x^m in (Z/NZ)[x]/(x^2 - P*x + Q), where x^m = U_m*x - Q*U_{m-1}.
The powers x^((N+1)/q) for all q are computed together by a product tree, at
the cost of about log2(n) full-size powers instead of n.

For each q the program uses the first P in 1, 3, 5, ... that satisfies the
conditions, and it tries D in the order 5, -7, 9, -11, ... This is the same
search order as the PARI/GP program A400659.gp, which was written separately
(different prime sieve, Jacobi symbol and ring arithmetic), so the two print
the same lines apart from the timings.

The terms found are compared with the b-file, which is embedded below.
"""

import sys
import time

try:
    from gmpy2 import mpz, gcd
except ImportError:          # plain Python integers work, only slower
    from math import gcd
    mpz = int

A400659 = [2, 1, 1, 3, 1, 1, 2, 5, 5, 4, 16, 4, 1, 12, 9, 2, 5, 3, 6, 2, 34, 3, 2, 1, 6, 65, 3, 5, 6, 3, 229, 11, 17, 5, 7, 16, 19, 15, 46, 11, 14, 20, 10, 207, 12, 12, 32, 63, 6, 93, 4, 16, 96, 5, 5, 329, 172, 34, 68, 4, 395, 150, 20, 2, 20, 1, 54, 1, 106, 5, 103, 11, 20, 34, 22, 63, 139, 48, 71, 25, 55, 45, 27, 22, 119, 9, 124, 470, 20, 51, 6, 71, 74, 95, 83, 589, 171, 30, 118, 1279]

SIEVE_LIMIT = 10 ** 5


def primes_below(limit):
    """All primes < limit, by the sieve of Eratosthenes."""
    flags = bytearray([1]) * limit
    flags[0:2] = b"\x00\x00"
    for p in range(2, int(limit ** 0.5) + 1):
        if flags[p]:
            flags[p * p::p] = bytearray(len(range(p * p, limit, p)))
    return [p for p in range(limit) if flags[p]]


SIEVE_PRIMES = primes_below(SIEVE_LIMIT)
PRIMES = SIEVE_PRIMES[:len(A400659)]
FERMAT_BASES = [p for p in SIEVE_PRIMES if p <= 97]


def jacobi(a, m):
    """Jacobi symbol (a/m) for odd m > 0, by quadratic reciprocity."""
    a %= m
    result = 1
    while a:
        while a % 2 == 0:
            a //= 2
            if m % 8 in (3, 5):
                result = -result
        a, m = m, a
        if a % 4 == 3 and m % 4 == 3:
            result = -result
        a %= m
    return result if m == 1 else 0


def fermat_witness(N):
    """A base b <= 97 with gcd(b, N) = 1 and b^(N-1) != 1 (mod N), or 0."""
    for b in FERMAT_BASES:
        if N % b and pow(mpz(b), N - 1, N) != 1:
            return b
    return 0


class Ring:
    """(Z/NZ)[x]/(x^2 - P*x + Q); the element u*x + v is the pair (u, v)."""

    X = (1, 0)

    def __init__(self, P, Q, N):
        self.P, self.Q, self.N = mpz(P), mpz(Q), N

    def square(self, s):
        a, b = s                                     # (a x + b)^2, with x^2 = P x - Q
        aa = a * a
        return ((aa * self.P + 2 * a * b) % self.N, (b * b - aa * self.Q) % self.N)

    def mul(self, s, t):
        (a, b), (c, d) = s, t
        ac, bd = a * c, b * d
        mid = (a + b) * (c + d) - ac - bd            # a d + b c, with one product
        return ((ac * self.P + mid) % self.N, (bd - ac * self.Q) % self.N)

    def mul_x(self, s):
        a, b = s                                     # (a x + b) x = (a P + b) x - a Q
        return ((a * self.P + b) % self.N, (-a * self.Q) % self.N)

    def power(self, s, e):
        """s^e for e >= 1; multiplying by x is free, so x^e is cheaper."""
        step = self.mul_x if s == Ring.X else (lambda r: self.mul(r, s))
        result = (mpz(s[0]), mpz(s[1]))
        for bit in bin(e)[3:]:
            result = self.square(result)
            if bit == "1":
                result = step(result)
        return result


def cofactor_powers(ring, base, qe):
    """[base^(F/q) for (q, e) in qe], where F = prod q^e, by a product tree."""
    if len(qe) == 1:
        q, e = qe[0]
        return [ring.power(base, q ** (e - 1)) if e > 1 else base]
    mid = len(qe) // 2
    left, right = qe[:mid], qe[mid:]
    f_left = f_right = 1
    for q, e in left:
        f_left *= q ** e
    for q, e in right:
        f_right *= q ** e
    return (cofactor_powers(ring, ring.power(base, f_right), left)
            + cofactor_powers(ring, ring.power(base, f_left), right))


def lucas_proof(N, n):
    """Lucas N+1 proof for N with N + 1 = A002110(n-1)*prime(n)^k.

    Returns (D, P, overrides) with overrides = [[q, P_q], ...] for the primes q
    whose P differs from the reported P, or None if no witness was found within
    the search limits (never a pass)."""
    F = N + 1
    if N < 3 or N % 2 == 0:
        return None
    qe, rest = [], F
    for q in PRIMES[:n]:
        e = 0
        while rest % q == 0:
            rest //= q
            e += 1
        qe.append((q, e))
    if rest != 1 or any(e == 0 for _, e in qe):
        raise ValueError("N+1 is not A002110(n-1)*prime(n)^k")
    # Selfridge order for D: 5, -7, 9, -11, 13, ...
    d = 5
    for t in range(1, 201):
        D = d if t % 2 else -d
        d += 2
        g = gcd(D, N)
        if g != 1:
            if g < N:
                return None
            continue
        if jacobi(D, N) == -1:
            break
        if t == 200:
            return None
    # For each q the first P = 1, 3, 5, ... that works; all still-open q are
    # handled together for each P.
    Ps = [0] * len(qe)
    pending = list(range(len(qe)))
    for P in range(1, 2002, 2):                      # D is odd, so P must be odd
        if not pending:
            break
        assert (P * P - D) % 4 == 0                 # every D tried is 1 mod 4
        Q = (P * P - D) // 4
        if Q == 0 or gcd(2 * Q, N) != 1:
            continue
        ring = Ring(P, Q, N)
        sub = [qe[i] for i in pending]
        f_sub = 1
        for q, e in sub:
            f_sub *= q ** e
        base = Ring.X if f_sub == F else ring.power(Ring.X, F // f_sub)
        zs = cofactor_powers(ring, base, sub)        # x^((N+1)/q) for each open q
        if ring.power(zs[0], sub[0][0])[0] != 0:     # U_{N+1} != 0: N is composite
            return None
        still = []
        for i, z in zip(pending, zs):
            if gcd(z[0], N) == 1:
                Ps[i] = P
            else:                                    # this P does not separate q
                still.append(i)
        pending = still
    if pending:
        return None
    # Report the most common P (the smallest, on a tie) and the exceptions.
    counts = {}
    for P in Ps:
        counts[P] = counts.get(P, 0) + 1
    P0 = max(sorted(counts), key=lambda P: counts[P])
    overrides = [[q, P] for (q, _), P in zip(qe, Ps) if P != P0]
    return D, P0, overrides


def prove_term(n):
    """Find and prove a(n). Prints one line; returns a(n), or -1 on failure."""
    t0 = time.time()
    p_n = PRIMES[n - 1]
    h = mpz(1)
    for p in PRIMES[:n - 1]:
        h *= p
    # Residues of h*p_n^k modulo each sieve prime p > p_n; p | N iff residue == 1.
    sp = [p for p in SIEVE_PRIMES if p > p_n]
    gs = [p_n % p for p in sp]
    rs = [int(h % p) for p in sp]
    k, H, one, by_factor, bases = 0, h, False, 0, set()
    while True:
        k += 1
        H *= p_n                                     # H = h * p_n^k = N + 1
        N = H - 1
        rs = [r * g % p for r, g, p in zip(rs, gs, sp)]
        if N == 1:                                   # n = 1, k = 1: 1 is not prime
            one = True
            continue
        if 1 in rs and N > sp[rs.index(1)]:          # a prime factor p < N, p < 10^5
            by_factor += 1
            continue
        b = fermat_witness(N)
        if b:
            bases.add(b)
            continue
        # No witness of compositeness: N must now be proved prime, or the run stops.
        w = lucas_proof(N, n)
        if w is None:
            print("n=%d k=%d: no compositeness witness and no Lucas proof; "
                  "check this number with a primality prover" % (n, k))
            return -1
        D, P0, overrides = w
        parts = []
        if one:
            parts.append("k=1 gives N=1 (not prime)")
        if by_factor or bases:
            how = []
            if by_factor:
                how.append("%d with a prime factor below 10^5" % by_factor)
            if bases:
                how.append("%d by a Fermat test, bases %s" % (k - 1 - one - by_factor, sorted(bases)))
            parts.append("k<%d composite (%s)" % (k, ", ".join(how)))
        below = "; ".join(parts) if parts else "no smaller k"
        print("n=%d  a(n)=%d  digits=%d  %s  |  k=%d prime: Lucas N+1, D=%d, P=%d%s  [%d ms]"
              % (n, k, len(str(N)), below, k, D, P0,
                 ", except (q,P)=%s" % overrides if overrides else "",
                 int((time.time() - t0) * 1000)), flush=True)
        return k


def verify(nmax=len(A400659)):
    if not 1 <= nmax <= len(A400659):
        raise SystemExit("the b-file covers n = 1..%d" % len(A400659))
    t0 = time.time()
    bad = 0
    for n in range(1, nmax + 1):
        r = prove_term(n)
        if r != A400659[n - 1]:
            bad += 1
            print("MISMATCH at n=%d: found %d, b-file has %d" % (n, r, A400659[n - 1]))
    if bad:
        print("FAILED: %d of %d terms not confirmed." % (bad, nmax))
    else:
        print("OK: a(1)..a(%d) proved and equal to the b-file. Total %d s."
              % (nmax, int(time.time() - t0)))
    return bad == 0


if __name__ == "__main__":
    sys.exit(0 if verify(int(sys.argv[1]) if len(sys.argv) > 1 else len(A400659)) else 1)