#!/usr/bin/env python3
# verify_baum_sweet_hankel.py
# Ancillary file for the note "The Hankel determinants of the Baum-Sweet sequence are 2-automatic".
# David Niedbala Giraudin, ORCID 0009-0009-1526-1178. Technical assistance: Claude (Anthropic).
#
# Plain Python 3, no dependency, exact arithmetic over F_2.  Run:  python3 verify_baum_sweet_hankel.py
# Polynomials over F_2 are Python integers (bit i = coefficient of t**i); addition is xor(a, b).
#
# Part 1  ground truth from first principles: Euclid's algorithm on the series, Gaussian elimination on
#         the Hankel matrices, Berlekamp-Massey.
# Part 2  the Mills-Robbins expansion Gamma and the identities of Section 3 of the note.
# Part 3  Lemma 4.1: the 20 factors of length 3 of Gamma (all finite data of the induction).
# Part 4  Theorem 5.1: the transducer for t/f (finite product with the factors of length 3).
# Part 5  Theorem 1.1: the 8-state automaton satisfies the recursion (R1)-(R4) and the base table
#         (finite reachability checks), hence computes chi(n+1) for every n.
# Part 6  the 2-kernels (74 and 68) and the minimal MSB-first automata of chi_0 and chi_1.
# Part 7  Corollary 1.3: the linear complexity profile.
# The script stops at the first failure and ends with "ALL CHECKS PASSED".

import sys
import time

T0 = time.time()


def check(cond, msg):
    if not cond:
        print("FAIL  " + msg, flush=True)
        sys.exit(1)
    print("  ok  %-86s [%6.1f s]" % (msg, time.time() - T0), flush=True)


def xor(a, b):
    return (a | b) & ~(a & b)


def deg(p):
    return p.bit_length() - 1


def pmul(a, b):
    r = 0
    while b:
        if b & 1:
            r = xor(r, a)
        a <<= 1
        b >>= 1
    return r


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


# ------------------------------------------------------------------------------------------------
print("Part 1. Ground truth", flush=True)


def baum_sweet(N):
    b = [0] * N
    b[0] = 1
    for n in range(1, N):
        if n % 2 == 1:
            b[n] = b[(n - 1) // 2]
        elif n % 4 == 0:
            b[n] = b[n // 4]
    return b


def bs_brute(n):
    if n == 0:
        return 1
    runs = [len(x) for x in format(n, 'b').split('1')]
    return 1 if all(r % 2 == 0 for r in runs) else 0


MM = 2 ** 17
b = baum_sweet(MM + 4)
check(all(bs_brute(n) == b[n] for n in range(1 << 14)), "recursion b_{2n+1}=b_n, b_{4n}=b_n, b_{4n+2}=0 == block definition (n < 2**14)")


def cf_quotients(shift, MMx):
    """partial quotients (after the integer part 0) of sum_{n>=1} b[n-1+shift] t**(-n), truncated at
    cumulative degree MMx//2 (enough precision: the series is known modulo t**(-MMx))"""
    num = 0
    for i in range(MMx):
        if b[i + shift]:
            num |= 1 << (MMx - 1 - i)
    r0, r1 = 1 << MMx, num
    qs = []
    cum = 0
    while r1:
        q, r = pdivmod(r0, r1)
        if cum + deg(q) > MMx // 2:
            break
        cum += deg(q)
        qs.append(q)
        r0, r1 = r1, r
    return qs


def cumdeg(qs):
    S = [0]
    for q in qs:
        S.append(S[-1] + deg(q))
    return S


q_alpha0 = cf_quotients(0, MM)      # alpha_0 = f/t = sum b_n t**(-n-1)  (Hankel determinants H_n of b)
q_frac = cf_quotients(1, MM)        # {f} = f + 1 = sum_{n>=1} b_n t**(-n) (shifted determinants)
LIM = MM // 2
S0 = set(cumdeg(q_alpha0))
S1 = set(cumdeg(q_frac))
print("      Euclid: %d quotients for alpha_0, %d for {f}" % (len(q_alpha0), len(q_frac)), flush=True)


def hankel_nonzero(shift, nmax):
    out = {0}
    for n in range(1, nmax + 1):
        rows = []
        for i in range(n):
            v = 0
            for j in range(n):
                if b[i + j + shift]:
                    v |= 1 << j
            rows.append(v)
        rank = 0
        for col in range(n):
            piv = None
            for r in range(rank, n):
                if rows[r] >> col & 1:
                    piv = r
                    break
            if piv is None:
                break
            rows[rank], rows[piv] = rows[piv], rows[rank]
            for r in range(n):
                if r != rank and rows[r] >> col & 1:
                    rows[r] = xor(rows[r], rows[rank])
            rank += 1
        if rank == n:
            out.add(n)
    return out


NDET = 300
check(hankel_nonzero(0, NDET) == {n for n in S0 if n <= NDET},
      "det(b_{i+j})_{i,j<n} != 0  <=>  n in Sigma_0 (Euclid on alpha_0), n <= %d" % NDET)
check(hankel_nonzero(1, NDET) == {n for n in S1 if n <= NDET},
      "det(b_{i+j+1})_{i,j<n} != 0  <=>  n in Sigma_1 (Euclid on {f}), n <= %d" % NDET)


def berlekamp_massey(s):
    """linear complexity profile L(s, N), N = 1..len(s), over F_2"""
    C = [1]
    B = [1]
    L = 0
    m = 1
    prof = []
    for n in range(len(s)):
        d = s[n]
        for i in range(1, L + 1):
            if i < len(C):
                d = (d + C[i] * s[n - i]) % 2
        if d == 0:
            m += 1
        else:
            Tm = C[:]
            Bs = [0] * m + B
            C = C + [0] * max(0, len(Bs) - len(C))
            for i in range(len(Bs)):
                C[i] = (C[i] + Bs[i]) % 2
            if 2 * L <= n:
                L = n + 1 - L
                B = Tm
                m = 1
            else:
                m += 1
        prof.append(L)
    return prof


NBM = 3000
LBM = berlekamp_massey(b[:NBM])

# ------------------------------------------------------------------------------------------------
print("Part 2. The Mills-Robbins expansion", flush=True)
X, X1, X2, X21 = 0, 1, 2, 3                       # letters X, X+1, X**2, X**2+1
POL = {X: 0b10, X1: 0b11, X2: 0b100, X21: 0b101}
DEG = {X: 1, X1: 1, X2: 2, X21: 2}
TAU = {X: X1, X1: X, X2: X21, X21: X2}


def Lam(n):
    if n % 2:
        k = (2 ** (n - 1) - 4) // 3
        return [X, X1, X2] + [X, X2] * k + [X1, X]
    k = (2 ** (n - 1) - 5) // 3
    return [X1, X21] + [X, X2] * k + [X, X21, X1]


NM = 19
Gam = {0: [X1], 1: [X1, X21], 2: [X1, X21, X, X2, X1]}
for n in range(3, NM + 1):
    r = list(reversed(Gam[n - 3]))
    r[0] = TAU[r[0]]
    Gam[n] = Gam[n - 1] + r + Lam(n) + Gam[n - 3]
nq = min(len(Gam[NM]), len(q_frac) - 1)
check(q_frac[0] == POL[X] and all(POL[Gam[NM][i]] == q_frac[i + 1] for i in range(nq)),
      "{f} = [0; X, Gamma] : Gamma == quotients of {f} by Euclid (%d letters)" % nq)
check(all(sum(DEG[c] for c in Gam[n]) == 2 ** (n + 1) - 1 for n in range(NM + 1)), "Lemma 3.1: D_n = 2**(n+1) - 1 (n <= %d)" % NM)
check(all(sum(DEG[c] for c in Lam(n)) == 2 ** (n - 1) + 2 and Lam(n) == Lam(n)[::-1] for n in range(3, NM + 1)),
      "Lambda_n palindromic, of total degree 2**(n-1) + 2")


def phi(w, mode):
    out = []
    for c in w:
        if DEG[c] == 1:
            out.append(1)
        else:
            out += [0, 1] if mode == 1 else [0, 0]
    return out


CHI = {}
for mode in (1, 0):
    A = {n: phi(Gam[n], mode) for n in Gam}
    Bw = {n: phi(Gam[n][::-1], mode) for n in Gam}
    Gr = {m: Bw[m - 3] + phi(Lam(m), mode) + A[m - 3] for m in range(5, NM + 1)}
    check(all(len(Gr[m]) == 2 ** m and A[m] == A[m - 1] + Gr[m] and Bw[m] == Gr[m] + Bw[m - 1] for m in range(5, NM + 1)),
          "phi_%d: |G_m| = 2**m, A_m = A_{m-1} G_m, B_m = G_m B_{m-1} (5 <= m <= %d)" % (mode, NM))
    CHI[mode] = [1, 1] + A[NM][:-1]
check(all((n in S1) == (CHI[1][n] == 1) for n in range(LIM)), "chi_1 = 11 . phi_1(Gamma)  (dictionary for {f}; n < %d)" % LIM)

# ------------------------------------------------------------------------------------------------
print("Part 3. Factors of length 3", flush=True)
fac = lambda w: {tuple(w[i:i + 3]) for i in range(len(w) - 2)}
F3 = fac(Gam[16])
check(len(F3) == 20 and {u[::-1] for u in F3} == F3, "|F| = 20 and F closed under reversal")
check(all(fac(Gam[n]) <= F3 for n in range(3, 9)), "base case: Fac_3(Gamma_n) subset F for 3 <= n <= 8")
check(all(fac(Lam(n)) <= F3 for n in range(3, 19)), "Fac_3(Lambda_n) subset F (n <= 18; constant in n for n >= 5)")
deb2 = {tuple(Gam[n][:2]) for n in range(1, NM + 1)}
fin3 = {}
for n in range(6, NM + 1):
    fin3.setdefault(n % 3, set()).add(tuple(Gam[n][-3:]))
check(len(deb2) == 1 and all(len(v) == 1 for v in fin3.values()), "begin_2(Gamma_n) constant, end_3(Gamma_n) depends on n mod 3")
FIN = {r: list(v)[0] for r, v in fin3.items()}
d2 = list(deb2)[0]
jok = True
for r6 in range(6):
    fp, fg = FIN[(r6 - 1) % 3], FIN[(r6 - 3) % 3]
    tg = (TAU[fg[2]], fg[1], fg[0])
    Ln = Lam(12 + r6 % 2)
    W = [(fp[1], fp[2], tg[0]), (fp[2], tg[0], tg[1]), tg, (d2[1], d2[0], Ln[0]), (d2[0], Ln[0], Ln[1]),
         (Ln[-2], Ln[-1], d2[0]), (Ln[-1], d2[0], d2[1])]
    jok = jok and set(W) <= F3
check(jok, "induction step: the 7 boundary factors of each class n mod 6 lie in F")
check(all(not (DEG[u[0]] == 2 and DEG[u[1]] == 2) for u in F3), "no two adjacent letters of degree 2")

# ------------------------------------------------------------------------------------------------
print("Part 4. The transducer for t/f", flush=True)
TT = 0b10


def consume(st, a):
    p, q, r, s = st
    return (xor(pmul(p, a), q), p, xor(pmul(r, a), s), r)


def emit(st, nxt):
    """emit c = polypart(p/r) while this is certain for every admissible next letter"""
    p, q, r, s = st
    out = []
    while r:
        c, rem = pdivmod(p, r)
        det = xor(pmul(p, s), pmul(q, r))
        ok = c != 0
        for e in nxt:
            ras = xor(pmul(r, POL[e]), s)
            if not (deg(ras) >= deg(r) and deg(det) < deg(r) + deg(ras)):
                ok = False
        if not ok:
            break
        out.append(c)
        p, q, r, s = r, s, rem, xor(q, pmul(c, s))
    return out, (p, q, r, s)


g = Gam[16]
st = (0, TT, 1, 0)                            # value t/f = (0 f + t)/(1 f + 0)
o_a, st = emit(consume(st, 1), [X])           # integer part 1 of f
o_b, st = emit(consume(st, POL[X]), [g[0]])   # then X
o_c, st = emit(consume(st, POL[g[0]]), [g[1]])
nx = [e for e in range(4) if (g[0], g[1], e) in F3]
o_d, st = emit(consume(st, POL[g[1]]), nx)
check(o_a == [] and o_b == [0b11] and o_c == [0b10] and o_d == [], "prefix: t/f = [X+1; X, ...], nothing emitted at Gamma[1] = X**2+1")
Dg = DEG[g[0]] + DEG[g[1]]
Eg = sum(deg(c) for c in o_b + o_c + o_d)
seen = set()
stack = [(st, (g[0], g[1]), Eg - Dg)]
seen.add(stack[0])
alph = set()
safe = True
book = True
blk = {X: 0b10, X1: 0b10, X2: 0b1000, X21: 0b1010}
while stack:
    st, (a, bb), dl = stack.pop()
    for c in range(4):
        if (a, bb, c) not in F3:
            continue
        nx = [e for e in range(4) if (bb, c, e) in F3]
        o, ns = emit(consume(st, POL[c]), nx)
        alph |= set(o)
        if xor(pmul(ns[0], ns[3]), pmul(ns[1], ns[2])) != TT:
            safe = False
        nd = dl - DEG[c] + sum(deg(x) for x in o)
        if DEG[c] == 1:
            book = book and len(o) == 1 and nd == 1 and o[0] == blk[bb]
        else:
            book = book and len(o) == 0 and nd == -1
        y = (ns, (bb, c), nd)
        if y not in seen:
            seen.add(y)
            stack.append(y)
check(Eg - Dg == -1, "after the prefix, (output degree E) - (degree D read in Gamma) = -1")
check(safe, "product transducer x F: %d configurations, determinant = t in all" % len(seen))
check(alph == {0b10, 0b1000, 0b1010}, "output alphabet {X, X**3, X**3+X}; every emission certain for all admissible tails")
check(book, "a degree-1 letter emits one quotient (X, X**3, X**3+X after deg 1, X**2, X**2+1), then E - D = 1;"
      " a degree-2 letter emits nothing, then E - D = -1")
check(all((n in S0) == (CHI[0][n] == 1) for n in range(LIM)), "hence chi_0 = 11 . phi_0(Gamma); checked against Euclid (n < %d)" % LIM)
check(all(CHI[0][n] == (CHI[1][n] & CHI[1][n - 1]) for n in range(1, len(CHI[0]))), "transfer relation chi_0(n) = chi_1(n) chi_1(n-1)")

# ------------------------------------------------------------------------------------------------
print("Part 5. The automaton", flush=True)
# omega_e(n) = chi_e(n+1); MSB-first automaton, state 0 initial, delta(0,0) = 0
DELTA = [(0, 1), (2, 3), (4, 5), (6, 0), (1, 4), (7, 5), (6, 7), (5, 6)]
OUT = {0: [1, 1, 0, 0, 1, 0, 0, 1], 1: [1, 1, 0, 1, 1, 0, 1, 1]}


def run(p, s):
    for ch in s:
        p = DELTA[p][int(ch)]
    return p


B4 = {mode: phi(Gam[4][::-1], mode) for mode in (0, 1)}
for mode in (1, 0):
    out = OUT[mode]
    om = CHI[mode][1:]

    def bfs(starts, step, test):
        seenb = set(starts)
        stk = list(starts)
        ok = True
        while stk:
            x = stk.pop()
            ok = ok and test(x)
            for c in (0, 1):
                y = step(x, c)
                if y not in seenb:
                    seenb.add(y)
                    stk.append(y)
        return ok

    def orbit(prefix, unit, a0):
        p = run(0, prefix + unit * a0)
        lst = []
        while p not in lst:
            lst.append(p)
            p = run(p, unit)
        return lst

    base = all(out[run(0, format(n, 'b'))] == om[n] for n in range(1, 256)) and out[0] == om[0]
    rA = bfs([(run(0, '111'), 0, 0, False)],
             lambda x, c: (DELTA[x[0]][c], DELTA[x[1]][c], min(x[2] + 1, 6), x[3] or c == 1),
             lambda x: not (x[2] >= 6 and x[3]) or out[x[0]] == out[x[1]])
    rB = all(bfs([(DELTA[p][0], run(0, '1'), 0)], lambda x, c: (DELTA[x[0]][c], DELTA[x[1]][c], min(x[2] + 1, 5)),
                 lambda x: x[2] < 5 or out[x[0]] == out[x[1]]) for p in orbit('100', '1', 0))
    rB4 = all(out[run(p, format(w, '05b'))] == B4[mode][w] for p in orbit('100', '1', 1) for w in range(31))
    d = (0, 1) if mode == 1 else (0, 0)
    core = (1, d[0], d[1])

    def expect(m_odd, i3, sp):
        if m_odd:
            return {'i1': 1, 'i2': d[0], 'i3': d[1], 'i2q': 1}.get(sp, core[(i3 - 4) % 3])
        return {'i1': d[0], 'i2': d[1], 'i2q-2': 1, 'i2q-1': d[0], 'i2q': d[1]}.get(sp, core[(i3 - 3) % 3])

    ZT = {('Z', 0): 'Z', ('Z', 1): 'Z1', ('Z1', 0): 'Z10'}
    NT = {('N', 1): 'N', ('N', 0): 'N0', ('N0', 1): 'N01'}
    r101 = bfs([(run(0, '101'), 0, 0, 0, 'Z')],
               lambda x, c: (DELTA[x[0]][c], (2 * x[1] + c) % 3, min(x[2] + 1, 6), (x[3] + 1) % 2, ZT.get((x[4], c), 'O')),
               lambda x: x[2] < 6 or out[x[0]] == expect(x[3] == 1, (x[1] + 1) % 3, {'Z': 'i1', 'Z1': 'i2', 'Z10': 'i3'}.get(x[4])))
    r110 = bfs([(run(0, '110'), 0, 0, 0, 'N')],
               lambda x, c: (DELTA[x[0]][c], (2 * x[1] + c) % 3, min(x[2] + 1, 6), (x[3] + 1) % 2, NT.get((x[4], c), 'O')),
               lambda x: x[2] < 6 or out[x[0]] == expect(x[3] == 1, ((1 if x[3] == 0 else 2) + x[1] + 1) % 3,
                                                        {'N': 'i2q', 'N0': 'i2q-1', 'N01': 'i2q-2'}.get(x[4])))
    rends = all(out[p] == 1 for p in orbit('100', '1', 6)) and all(out[p] == 1 for p in orbit('111', '0', 6))
    check(base and DELTA[0][0] == 0, "e=%d: base table n < 256 and leading-zero invariance" % mode)
    check(rA and rB and rB4, "e=%d: rules (R1) omega(111y) = omega(y), (R2) omega(100 1..1 0 z) = omega(1z), (R3) B_4 block" % mode)
    check(rends and r101 and r110, "e=%d: rule (R4): the block L_m (both parities, all positions)" % mode)
    check(all(out[run(0, format(n, 'b'))] == om[n] for n in range(1, len(om))),
          "e=%d: automaton == chi_%d(n+1) on the whole prefix (n < %d)" % (mode, mode, len(om)))
    Sx = S0 if mode == 0 else S1
    check(all((out[run(0, format(n - 1, 'b'))] == 1) == (n in Sx) for n in range(1, LIM)),
          "e=%d: automaton == Euclid's Sigma_%d directly (1 <= n < %d)" % (mode, mode, LIM))

# ------------------------------------------------------------------------------------------------
print("Part 6. Kernels and minimal automata of chi", flush=True)
for mode in (0, 1):
    out = OUT[mode]
    start = ('Z',)

    def step(stt, c):
        if stt == start:
            return start if c == 0 else (DELTA[0][1], 0)
        s_, dd = stt
        return (DELTA[s_][1], DELTA[s_][0]) if c == 1 else (DELTA[s_][0], DELTA[dd][1])
    Q = [start]
    idx = {start: 0}
    TR = []
    i = 0
    while i < len(Q):
        row = []
        for c in (0, 1):
            y = step(Q[i], c)
            if y not in idx:
                idx[y] = len(Q)
                Q.append(y)
            row.append(idx[y])
        TR.append(row)
        i += 1
    O = [1 if s_ == start else out[s_[1]] for s_ in Q]
    part = O[:]
    while True:
        sg = [(part[i], part[TR[i][0]], part[TR[i][1]]) for i in range(len(Q))]
        ids = {}
        new = [ids.setdefault(x, len(ids)) for x in sg]
        if len(set(new)) == len(set(part)):
            break
        part = new
    nmin = len(set(part))
    vec0 = tuple(range(len(Q)))
    seenv = {vec0}
    stk = [vec0]
    maps = set()
    while stk:
        v = stk.pop()
        maps.add(tuple(O[p] for p in v))
        for c in (0, 1):
            nv = tuple(TR[p][c] for p in v)
            if nv not in seenv:
                seenv.add(nv)
                stk.append(nv)
    check(nmin == (15 if mode == 0 else 8), "chi_%d: minimal MSB-first automaton has %d states" % (mode, nmin))
    check(len(maps) == (74 if mode == 0 else 68), "chi_%d: 2-kernel has %d elements (= states of the minimal LSB-first automaton)" % (mode, len(maps)))

# ------------------------------------------------------------------------------------------------
print("Part 7. Linear complexity profile", flush=True)
S0s = sorted(S0)


def L_from_sigma(N):
    j = 0
    while not (S0s[j] + S0s[j + 1] <= N < S0s[j + 1] + S0s[j + 2]):
        j += 1
    return S0s[j + 1]


check(all(LBM[N - 1] == L_from_sigma(N) for N in range(1, NBM + 1)),
      "Berlekamp-Massey L(b,N) == deg Q_j with deg Q_{j-1}+deg Q_j <= N < deg Q_j+deg Q_{j+1} (N <= %d)" % NBM)
gaps = {S0s[i + 1] - S0s[i] for i in range(len(S0s) - 1)}
check(gaps == {1, 3}, "gaps of Sigma_0 are 1 and 3 (up to %d)" % S0s[-1])
devs = {2 * L_from_sigma(N) - N for N in range(1, S0s[-3])}
check(devs == {-2, -1, 0, 1, 2, 3}, "2L(b,N) - N takes exactly the values -2..3 (N < %d)" % S0s[-3])
print("ALL CHECKS PASSED", flush=True)