"""Linked pairs of characters on Z_N (Section 2), the counting bound of Lemma 3.1 and Theorem 3.4,
prime powers and coprime products.   Output: linked.log"""
import sys
from schur import *

def linked(N, E, y, z):
    """y, z, y+z have trivial stabilizers in E, and ay+bz = c(y+z) only for a=b (Definition 2.2)."""
    k = len(E)
    for x in (y, z, y + z):
        if len({a * x % N for a in E}) < k: return False
    orb = {c * (y + z) % N for c in E}
    return all(a == b or (a * y + b * z) % N not in orb for a in E for b in E)

def two_maximal(N, E):
    d2 = {(2 * (c - d)) % N for c in E for d in E}
    return all(b in (1, N - 1) or (1 - b) % N not in d2 for b in E)

def partners(N, E, y=1):
    return [z for z in range(1, N) if linked(N, E, y, z)]

if __name__ == "__main__":
    out = open("linked.log", "w")
    def say(*a):
        print(*a); print(*a, file=out); out.flush()
    P = int(sys.argv[1]) if len(sys.argv) > 1 else 400
    say("PRIME CASE, p < %d, even k | p-1, k < p-1." % P)
    say("  LINK = some z is linked to 1 (Theorem 2.4 applies);  QUAD = p > (k-1)(k-2)+2 (Theorem 3.4);  BBC = E is 2-maximal")
    tot = lk = quad = bbc = 0; worst = 0
    for p in primes(P):
        if p < 5: continue
        for k in range(2, p - 1, 2):
            if (p - 1) % k: continue
            E = cyclic_subgroup(p, k)
            tot += 1
            good = [z for z in range(1, p - 1) if linked(p, E, 1, z)]
            bad = (p - 2) - len(good)                      # z != 0, -1 not linked to 1
            assert bad <= (k - 1) * (k - 2), (p, k)        # Lemma 3.1
            if k > 2: worst = max(worst, bad / ((k - 1) * (k - 2)))
            q = p > (k - 1) * (k - 2) + 2
            assert good or not q                           # Theorem 3.4
            lk += bool(good); quad += q; bbc += two_maximal(p, E)
            if p < 62: say("  p=%2d k=%2d  LINK=%d QUAD=%d BBC=%d" % (p, k, bool(good), q, two_maximal(p, E)))
    say("TOTAL pairs (p,k): %d;  LINK: %d;  QUAD: %d;  BBC: %d;  Lemma 3.1 checked, max bad/bound = %.3f" % (tot, lk, quad, bbc, worst))
    say("")
    say("PRIME POWERS p^n, n >= 2, E of order k | p-1: is p linked to 1 ?")
    for p in (3, 5, 7, 11, 13, 17, 19, 23):
        for n in (2, 3):
            N = p ** n
            if N > 3000: continue
            for k in range(2, p, 2):
                if (p - 1) % k: continue
                say("  N=%4d k=%2d  linked(1,p)=%s" % (N, k, linked(N, cyclic_subgroup(N, k), 1, p)))
    say("")
    say("COPRIME PRODUCTS N = N1*N2, E injective mod N1 and mod N2: are e1, e2 linked ?")
    for N1, N2 in [(5, 13), (7, 13), (5, 17), (13, 17)]:
        N = N1 * N2
        e1 = next(x for x in range(N) if x % N1 == 1 and x % N2 == 0)
        e2 = next(x for x in range(N) if x % N1 == 0 and x % N2 == 1)
        done = set()
        for a in range(2, N):
            if math.gcd(a, N) > 1: continue
            E = [1]; x = a
            while x != 1: E.append(x); x = x * a % N
            E = tuple(sorted(E))
            if E in done or N - 1 not in E or len(E) == 2: continue
            done.add(E)
            if len({e % N1 for e in E}) == len(E) and len({e % N2 for e in E}) == len(E):
                say("  N=%d E=%s  linked(e1,e2)=%s" % (N, list(E), linked(N, E, e1, e2)))
    say("")
    say("GENERAL CASE (Lemma 3.1, Theorems 3.2 and 3.4): N < 130, all the groups -1 in E generated by at most 3 elements,")
    say("  m = max gcd(a-1,N) over a != 1 in E.  Checks: non-regular characters <= (k-1)(m-1)+1, and for any regular y,")
    say("  the number of z not linked to y is <= mk(k-1)-2(k-2).")
    groups = 0; worst = 0; applies = 0; fixed = 0
    for N in range(5, 130):
        units = [a for a in range(1, N) if math.gcd(a, N) == 1]
        done = set()
        for a in units:
            for b in units:
                if b < a: continue
                E = {1, N - 1}
                while True:
                    F = E | {x * g % N for x in E for g in (a, b)}
                    if F == E: break
                    E = F
                E = tuple(sorted(E))
                if E in done or len(E) > 12: continue
                done.add(E); groups += 1
                k = len(E); m = max(math.gcd(c - 1, N) for c in E if c != 1)
                reg = [y for y in range(N) if len({c * y % N for c in E}) == k]
                assert N - len(reg) <= (k - 1) * (m - 1) + 1, (N, E)
                bound = m * k * (k - 1) - 2 * (k - 2)
                for y in reg:
                    bad = sum(1 for z in range(N) if not linked(N, E, y, z))
                    assert bad <= bound, (N, E, y)
                    worst = max(worst, bad / bound)
                if N > bound:
                    applies += 1; fixed += m > 1
                    assert any(linked(N, E, 1, z) for z in range(N)), (N, E)
    say("  groups checked: %d, max bad/bound = %.3f, Theorem 3.4 applies to %d of them, among which %d with m > 1" % (
        groups, worst, applies, fixed))
    out.close()