#!/usr/bin/env python3
# verify_collatz_binary.py
# Ancillary file for the note "The binary Collatz dynamics of the trinomials x^n+x+1:
# proof of two conjectures of Gallardo and Rahavandrainy".
# David Niedbala Giraudin, ORCID 0009-0009-1526-1178. Technical assistance: Claude (Anthropic).
#
# Plain Python 3, no dependency, no input file. Run:  python3 verify_collatz_binary.py
# Exact arithmetic only. A polynomial over F_2 is an integer whose bit i is the coefficient
# of x**i. D = x(x+1) = 6, M = x^2+x+1 = 7, M + D = 1.
# The Collatz step of Gallardo-Rahavandrainy on odd A is  T(A) = odd part of 1 + M*A.
# The script stops at the first failure and ends with "ALL CHECKS PASSED".

import math
import sys

D_POLY = 0b110
M_POLY = 0b111


def deg(a):
    return a.bit_length() - 1 if a else -1


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


def ppow(a, e):
    r = 1
    for _ in range(e):
        r = pmul(r, a)
    return r


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


def vx(a):
    v = 0
    while not (a >> v) & 1:
        v += 1
    return v


def vx1(a):
    v = 0
    while True:
        q, r = pdivmod(a, 0b11)
        if r:
            return v
        a = q
        v += 1


def oddpart(a):
    while a % 2 == 0:
        a >>= 1
    while True:
        q, r = pdivmod(a, 0b11)
        if r:
            return a
        a = q


def step(u):
    return oddpart(1 ^ pmul(M_POLY, u))


def mu(u):
    z = u ^ 1
    return min(vx(z), vx1(z))


def orbit(u):
    seq = [u]
    while u != 1:
        u = step(u)
        seq.append(u)
    return seq


def T(n):
    return (1 << n) | 0b11


def F(m):
    return pmul(1 << m, M_POLY) ^ 1


def V(m):
    return oddpart(F(m))


def k_of(n):
    return math.ceil(math.log2(n))


def mprime(m):
    if m % 2 == 0:
        return m // 2 - 1
    if m % 4 == 1:
        return (m - 3) // 2
    return (m - 1) // 2


def W_of(u):
    m = mu(u)
    acc = u ^ 1
    for _ in range(m):
        q, r = pdivmod(acc, D_POLY)
        assert r == 0
        acc = q
    return pmul(acc, ppow(M_POLY, m))


def check(cond, msg):
    if not cond:
        print("FAILED: " + msg, flush=True)
        sys.exit(1)
    print("  ok  " + msg, flush=True)


print("Binary Collatz dynamics of x^n+x+1: verification", flush=True)

# identity (1) and the three identities used in Lemma 8
X1 = 0b11
check(M_POLY ^ D_POLY == 1, "M + D = 1")
check(M_POLY ^ X1 == 0b100 and M_POLY ^ pmul(X1, X1) == 0b10
      and M_POLY ^ ppow(X1, 3) == 0b1000, "M+(x+1) = x^2, M+(x+1)^2 = x, M+(x+1)^3 = x^3")

# Lemma 3 (plateau lemma) along every orbit, n = 3..199
for n in range(3, 200):
    seq = orbit(T(n))
    for i in range(len(seq) - 1):
        u, w = seq[i], seq[i + 1]
        same = deg(w) == deg(u)
        assert deg(w) <= deg(u), (n, i)
        assert same == (mu(u) >= 2), (n, i)
        if same:
            assert w ^ 1 == pmul(M_POLY, pdivmod(u ^ 1, D_POLY)[0]) and mu(w) == mu(u) - 1
check(True, "Lemma 3: degree kept iff mu >= 2, T(u)+1 = M z / D, mu drops by 1 (n = 3..199)")

# Lemma 4 (plateau-to-plateau map and stopping criterion) along every orbit, n = 3..159
for n in range(3, 160):
    seq = orbit(T(n))
    i = 0
    while i < len(seq) - 1:
        u = seq[i]
        m = mu(u)
        W = W_of(u)
        assert deg(W) == deg(u)
        assert seq[i + m] == oddpart(W ^ 1)
        assert (seq[i + m] == 1) == (oddpart(W ^ 1) == 1)
        i += m
check(True, "Lemma 4: deg W = deg u, T^mu(u) = odd(W+1), stop iff W+1 = x^i(x+1)^j (n = 3..159)")

# Proposition 5 (single-plateau orbits), all admissible 0 <= i, j <= 8, e = 0..3
count = 0
for i in range(9):
    for j in range(9):
        if (i, j) == (0, 0) or (i - j) % 3:
            continue
        E = pmul(1 << i, ppow(X1, j)) ^ 1
        q, r = pdivmod(E, M_POLY)
        assert r == 0
        P = 1 ^ pmul(D_POLY, q)
        assert deg(P) == i + j and P & 1 and oddpart(P) == P
        for e in range(4):
            u = ppow(P, 2 ** e)
            seq = orbit(u)
            assert len(seq) == 2 ** e + 1 and all(deg(t) == 2 ** e * (i + j) for t in seq[:-1])
        count += 1
check(count == 26, "Proposition 5: the 26 pairs (i, j), e = 0..3, give 2^e terms then 1")
check(ppow(1, 1) == 1 and T(3) == 1 ^ pmul(D_POLY, pdivmod(pmul(1 << 3, 1) ^ 1, M_POLY)[0])
      and T(4) == 1 ^ pmul(D_POLY, pdivmod(pmul(1 << 2, ppow(X1, 2)) ^ 1, M_POLY)[0]),
      "T_3 = P_{3,0} and T_4 = P_{2,2}")

# Lemma 6 (first step)
check(all(mu(T(n)) == 1 and step(T(n)) == V(n - 3) for n in range(3, 400)),
      "Lemma 6: mu(T_n) = 1 and T(T_n) = V(n-3) (n = 3..399)")

# Lemma 7 (valuations)
expected = {0: 1, 1: 3, 2: 1, 3: 2}
check(all(vx(F(m)) == 0 and vx1(F(m)) == expected[m % 4]
          and deg(V(m)) == m + 2 - expected[m % 4] for m in range(1, 3000)),
      "Lemma 7: v_x(F_m) = 0, v_{x+1}(F_m) = 1,3,1,2, deg V(m) (m = 1..2999)")
check(V(0) == 1 and V(1) == 1 and all(deg(V(m)) >= 3 for m in range(2, 3000)),
      "Lemma 7: V(0) = V(1) = 1 and deg V(m) >= 3 for m = 2..2999")

# Lemma 8 (renormalisation)
check(all(step(V(m)) == pmul(V(mprime(m)), V(mprime(m))) and mu(V(m)) == 1
          for m in range(2, 3000)),
      "Lemma 8: T(V(m)) = V(m')^2 and mu(V(m)) = 1 (m = 2..2999)")

# Proposition 9 (Frobenius induction)
for m in range(2, 200):
    for e in range(4):
        u = ppow(V(m), 2 ** e)
        assert mu(u) == 2 ** e
        w = u
        for _ in range(2 ** e):
            w = step(w)
        assert w == ppow(V(mprime(m)), 2 ** (e + 1))
check(True, "Proposition 9: mu(V(m)^(2^e)) = 2^e and T^(2^e) gives V(m')^(2^(e+1)) "
            "(m = 2..199, e = 0..3)")

# Lemma 10 (binary recursion)
check(all(mprime(N - 2) + 2 == 2 * (N >> 2) + ((N & 1) | ((N >> 1) & 1))
          and (2 * (N >> 2) + ((N & 1) | ((N >> 1) & 1))).bit_length() == N.bit_length() - 1
          for N in range(4, 200000)),
      "Lemma 10: N' = 2 floor(N/4) + (b0 or b1), one binary digit less (N = 4..199999)")

# Theorem 1: length
check(all(len(orbit(T(n))) == 2 ** (k_of(n) - 2) + 1 for n in range(3, 1200)),
      "Theorem 1: length 2^(k-2)+1, k = ceil(log2 n) (n = 3..1199)")

# Theorem 1: blocks, and the starting points of the blocks
for n in range(5, 700):
    seq = orbit(T(n))
    k = k_of(n)
    i, j, m = 1, 1, n - 3
    degs = [deg(seq[0])]
    while seq[i] != 1:
        assert seq[i] == ppow(V(m), 2 ** (j - 1))
        L = 2 ** (j - 1)
        assert all(deg(t) == deg(seq[i]) for t in seq[i:i + L])
        degs.append(deg(seq[i]))
        i += L
        j += 1
        m = mprime(m)
    assert j - 1 == k - 2 and all(a > b for a, b in zip(degs, degs[1:]))
check(True, "Theorem 1: blocks B_j start at V(m_j)^(2^(j-1)), have 2^(j-1) terms, "
            "degrees strictly decrease, k-2 blocks (n = 5..699)")

# Remark: the last block
SEEDS = {4: 0b1011, 5: 0b1101, 6: 0b101111, 7: 0b10011}
for n in range(5, 700):
    k = k_of(n)
    b = n - 1
    top = (b >> (k - 2)) & 1
    low = 1 if b & ((1 << (k - 2)) - 1) else 0
    Nlast = 4 + 2 * top + low
    m = n - 3
    for _ in range(k - 3):
        m = mprime(m)
    assert m + 2 == Nlast and V(m) == SEEDS[Nlast]
check(True, "Remark: last block from V(2), V(3), V(4), V(5) by the three-digit rule (n = 5..699)")

# Section 4: the restatement with a non-strict inequality


def s_strict(n):
    s = 0
    while n - 2 ** (s + 2) >= 1:
        s += 1
    return s


def s_weak(n):
    s = 1
    while n - 2 ** (s + 2) >= 0:
        s += 1
    return s


truth = {n: 2 ** (k_of(n) - 2) + 1 for n in range(4, 2100)}
bad_strict = [n for n in truth if 2 ** s_strict(n) + 1 != truth[n]]
bad_weak = [n for n in truth if 2 ** s_weak(n) + 1 != truth[n]]
check(bad_strict == [], "Section 4: the statement with n - 2^(s+1) >= 1 holds (n = 4..2099)")
check(bad_weak == [4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048],
      "Section 4: the statement with n - 2^(s+1) >= 0 fails exactly at 4, 8, ..., 2048 (n = 4..2099)")
check(truth[32] == 9 and 2 ** s_weak(32) + 1 == 17, "Section 4: n = 32, length 9, restated formula 17")

print("ALL CHECKS PASSED", flush=True)