#!/usr/bin/env python3
"""Checks every identity and number in contraction_note.tex.

Usage: python3 verify_note.py [TMAX]
  TMAX (default 3) is the largest t for which the brute-force subtree enumeration is run
  (t = 6 takes about an hour in pure Python; t <= 3 takes seconds).
Exact integer / rational arithmetic throughout. Exits non-zero on any failed check.
"""
import sys
from fractions import Fraction as Fr
from math import comb
from itertools import combinations

fails = 0
def check(name, ok):
    global fails
    print(("ok   " if ok else "FAIL ") + name)
    if not ok: fails += 1

# ---------- the four polynomials of the note ----------
def N_(t): return 20*t**4 + 68*t**3 + 73*t**2 + t
def R_(t): return Fr(428*t**5 + 1288*t**4 + 1085*t**3 + 119*t**2 + 2*t, 3)
def A_(t): return 19*t**4 + 54*t**3 + 40*t**2 + 4*t
def B_(t): return Fr(404*t**5 + 1096*t**4 + 807*t**3 + 116*t**2 + 4*t, 3)
def NG(t): return N_(t) + A_(t) + 1
def RG(t): return R_(t) + B_(t) + A_(t) + 1
def diff(t): return RG(t)/NG(t) - R_(t)/N_(t)

# ---------- 1. cycle-path table by direct enumeration ----------
def table(t):
    m = 8*t; hubs = {0, 2*t}; v = t; out = {}
    for start in range(m):
        for L in range(1, m+1):
            vs = {(start+i) % m for i in range(L)}
            key = (len(vs & hubs), v in vs)
            c, o = out.get(key, (0, 0)); out[key] = (c+1, o+L)
    return out
ok = True
for t in range(1, 7):
    tb = table(t)
    allp = {h: tuple(map(sum, zip(tb.get((h, False), (0, 0)), tb.get((h, True), (0, 0))))) for h in range(3)}
    thr = {h: tb.get((h, True), (0, 0)) for h in range(3)}
    ok &= allp[0] == (20*t*t-4*t, Fr(112*t**3-4*t, 3)) and allp[1] == (24*t*t, 96*t**3) \
        and allp[2] == (20*t*t+4*t, Fr(368*t**3+96*t*t+4*t, 3))
    ok &= thr[0] == (t*t, t**3) and thr[1] == (12*t*t, 54*t**3) and thr[2] == (19*t*t+4*t, Fr(347*t**3+96*t*t+4*t, 3))
check("cycle-path table (both one-hub rows together) matches enumeration, t = 1..6", ok)

# ---------- 2. gap lemma Q, S against enumeration ----------
def QS_lemma(m, gaps):
    if not gaps: return m*m, Fr(m*m*(m+1), 2)
    Q = m + sum(comb(d, 2) for d in gaps); return Q, Fr(m*Q - sum(comb(d+1, 3) for d in gaps))
def QS_enum(m, marks):
    Q = S = 0
    for start in range(m):
        for L in range(1, m+1):
            vs = {(start+i) % m for i in range(L)}
            if marks <= vs: Q += 1; S += L
    return Q, Fr(S)
def gaps(marks, m):
    ms = sorted(marks)
    return [((ms[(i+1) % len(ms)] - ms[i]) % m) or m for i in range(len(ms))] if ms else []
ok = True
for t in range(1, 6):
    m = 8*t
    for S_ in [set(), {0}, {2*t}, {0, 2*t}, {t}, {0, t}, {t, 2*t}, {0, t, 2*t}]:
        ok &= QS_lemma(m, gaps(S_, m)) == QS_enum(m, S_)
check("gap lemma Q, S match enumeration for all mark sets, t = 1..5", ok)

# ---------- 3. table formulas give the four polynomials ----------
ok = True
for t in range(1, 25):
    N = (20*t*t-4*t) + 2*12*t*t*(t+1) + (20*t*t+4*t)*(t+1)**2 + t*(t+1)
    R = Fr(112*t**3-4*t, 3) + 2*(t+1)*(48*t**3 + 6*t**3) + (t+1)**2*(Fr(368*t**3+96*t*t+4*t, 3) + (20*t*t+4*t)*t) + Fr(t*(t+1)*(t+2), 3)
    A = t*t + 2*6*t*t*(t+1) + (19*t*t+4*t)*(t+1)**2
    B = t**3 + 2*(t+1)*(27*t**3 + 3*t**3) + (t+1)**2*(Fr(347*t**3+96*t*t+4*t, 3) + (19*t*t+4*t)*t)
    ok &= (N, R, A, B) == (N_(t), R_(t), A_(t), B_(t))
check("weighted table sums equal the four polynomials, t = 1..24", ok)

# ---------- 3b. equation (2) of the note: direct subset sums with the lemma values ----------
def eq2(t, with_v):
    m = 8*t; totN = 0; totR = Fr(0)
    for r in range(3):
        for X in combinations([0, 2*t], r):
            marks = set(X) | ({t} if with_v else set())
            Q, S = QS_lemma(m, gaps(marks, m))
            totN += t**r * Q; totR += t**r * (S + Fr(t+1, 2) * r * Q)
    if not with_v: totN += t*(t+1); totR += Fr(t*(t+1)*(t+2), 3)
    return totN, totR
check("equation (2) with lemma values gives N_t, R_t, A_t, B_t for t = 1..29",
      all(eq2(t, False) == (N_(t), R_(t)) and eq2(t, True) == (A_(t), B_(t)) for t in range(1, 30)))

# ---------- 4. F/(3DE), the expansion, and the threshold certificate ----------
def padd(a, b):
    n = max(len(a), len(b)); return [(a[i] if i < len(a) else 0) + (b[i] if i < len(b) else 0) for i in range(n)]
def pmul(a, b):
    r = [0]*(len(a)+len(b)-1)
    for i, x in enumerate(a):
        for j, y in enumerate(b): r[i+j] += x*y
    return r
def shift(p, a):
    res, pw = [0], [1]
    for c in p: res = padd(res, [x*c for x in pw]); pw = pmul(pw, [a, 1])
    return res
def ev(p, t): return sum(c*t**i for i, c in enumerate(p))
Np = [0, 1, 73, 68, 20]; R3 = [0, 2, 119, 1085, 1288, 428]; Ap = [0, 4, 40, 54, 19]; B3 = [0, 4, 116, 807, 1096, 404]
NGp = padd(padd(Np, Ap), [1]); RG3 = padd(padd(padd(R3, B3), [3*x for x in Ap]), [3])
D = [1, 73, 68, 20]; E = [1, 5, 113, 122, 39]
F = [1, 108, -33, 8849, 32814, 41102, 19989, 2948, -52]
num = padd(pmul(RG3, Np), [-x for x in pmul(R3, NGp)])      # 3 (R(G) N - R N(G))
while num[-1] == 0: num.pop()
check("N_t = t D and N(G_t) = E", Np == [0] + D and NGp == E)
check("mu(G_t) - mu(H_t) = F / (3 D E) as polynomials", num == [0] + F)
check("D and E have positive coefficients", all(c > 0 for c in D) and all(c > 0 for c in E))
c1 = Fr(F[8], 3*D[3]*E[4])
check("leading difference -t/45", c1 == Fr(-1, 45))
Q = padd(F, [-x for x in pmul(D, E)])
for name, poly in (("F", F), ("F-DE", Q)):
    translated = shift(poly, 1)
    check(f"{name}(1+s) has one negative leading coefficient and all others positive",
          translated[-1] < 0 and all(c > 0 for c in translated[:-1]))
check("threshold brackets F(62)>0>F(63) and (F-DE)(47)>0>(F-DE)(48)",
      ev(F, 62) > 0 > ev(F, 63) and ev(Q, 47) > 0 > ev(Q, 48))
F1 = shift(F, 1); Q1 = shift(Q, 1)
check("Descartes: F(1+s) has positive coefficients except a negative leading one", all(c > 0 for c in F1[:-1]) and F1[-1] < 0)
check("Descartes: (F-DE)(1+s) has positive coefficients except a negative leading one", all(c > 0 for c in Q1[:-1]) and Q1[-1] < 0)
check("F(1) > 0, F(62) > 0 > F(63); (F-DE)(47) > 0 > (F-DE)(48)", ev(F, 1) > 0 and ev(F, 62) > 0 > ev(F, 63) and ev(Q, 47) > 0 > ev(Q, 48))
ok = True
for t in range(1, 9):
    m = 8*t
    ok &= QS_lemma(m, [8*t]) == (32*t*t+4*t, Fr(512*t**3+96*t*t+4*t, 3)) and QS_lemma(m, [2*t, 6*t]) == (20*t*t+4*t, Fr(368*t**3+96*t*t+4*t, 3))
    ok &= QS_lemma(m, [t, 7*t]) == (25*t*t+4*t, Fr(428*t**3+96*t*t+4*t, 3)) and QS_lemma(m, [t, t, 6*t]) == (19*t*t+4*t, Fr(347*t**3+96*t*t+4*t, 3))
    ok &= QS_lemma(m, []) == (64*t*t, 256*t**3+32*t*t)
check("(Q,S) closed forms quoted in the proof for the five mark sets", ok)
check("difference < 1/3 first at t = 48 and < 0 first at t = 63 (direct, t <= 70)",
      [t for t in range(1, 71) if diff(t) < Fr(1, 3)][0] == 48 and [t for t in range(1, 71) if diff(t) < 0][0] == 63)

# ---------- 5. certificate at t = 63 ----------
check("t = 63 table", (N_(63), R_(63), NG(63), RG(63)) == (332352216, 148441475868, 645320026, 288225062749))
check("t = 63 fraction", diff(63) == Fr(-4293783638, 49101085318287))
check("exact local-mean leading coefficient -13/285",
      Fr(B3[-1], 3*Ap[-1]) - Fr(R3[-1], 3*Np[-1]) == Fr(-13, 285))
check("exact limiting rooted weight 19/39", Fr(Ap[-1], NGp[-1]) == Fr(19, 39))
t = 63
check("remark identity at t = 63", diff(t) == (A_(t)*(B_(t)/A_(t) - R_(t)/N_(t) + 1) + 1 - R_(t)/N_(t)) / NG(t))

# ---------- 6. brute force: connected vertex sets weighted by spanning-tree counts ----------
def build(t, leaf):
    m = 8*t; edges = [(i, (i+1) % m) for i in range(m)]; nxt = m
    for h in (0, 2*t):
        prev = h
        for _ in range(t): edges.append((prev, nxt)); prev = nxt; nxt += 1
    if leaf: edges.append((t, nxt)); nxt += 1
    adj = [set() for _ in range(nxt)]
    for a, b in edges: adj[a].add(b); adj[b].add(a)
    return adj
def det_int(M):                      # Bareiss, exact
    M = [row[:] for row in M]; n = len(M); sign = 1; prev = 1
    for k in range(n-1):
        if M[k][k] == 0:
            for i in range(k+1, n):
                if M[i][k] != 0: M[k], M[i] = M[i], M[k]; sign = -sign; break
            else: return 0
        for i in range(k+1, n):
            for j in range(k+1, n): M[i][j] = (M[i][j]*M[k][k] - M[i][k]*M[k][j]) // prev
        prev = M[k][k]
    return sign*M[n-1][n-1]
def tau(adj, S):                     # number of spanning trees of the induced subgraph
    S = sorted(S)
    if len(S) == 1: return 1
    idx = {v: i for i, v in enumerate(S)}; n = len(S); L = [[0]*n for _ in range(n)]
    for a in S:
        for b in adj[a]:
            if b in idx: L[idx[a]][idx[a]] += 1; L[idx[a]][idx[b]] -= 1
    return det_int([row[1:] for row in L[1:]])
def connected_sets(adj):             # ESU enumeration, each connected vertex set once
    n = len(adj)
    def ext(sub, cand, root):
        yield sub
        cand = list(cand)
        while cand:
            w = cand.pop()
            nb = {u for u in adj[w] if u > root and u not in sub and not any(u in adj[s] for s in sub)}
            yield from ext(sub | {w}, set(cand) | nb, root)
    for r in range(n):
        yield from ext(frozenset([r]), {u for u in adj[r] if u > r}, r)
def brute(adj):
    N = R = 0
    for S in connected_sets(adj):
        k = tau(adj, S); N += k; R += k*len(S)
    return N, R
tmax = int(sys.argv[1]) if len(sys.argv) > 1 else 3
for t in range(1, tmax+1):
    bh, bg = brute(build(t, False)), brute(build(t, True))
    check(f"brute force t = {t}: H {bh}, G {bg} equal the polynomials",
          bh == (N_(t), R_(t)) and bg == (NG(t), RG(t)))

print("\nALL CHECKS PASSED" if fails == 0 else f"\n{fails} CHECK(S) FAILED")
sys.exit(1 if fails else 0)