#!/usr/bin/env python3
"""Exact checks for the Pell/PTE family used on 4 < lambda < 5.

Only Python's standard library is required.  All checks use integers or
fractions.Fraction; no floating-point arithmetic is used.  The family theorem
and its limiting energy are proved in the paper.  This script checks the first
six nontrivial pairs, the exact fourth-moment formula, the first two monic
quotients, and the rational limiting energy at alpha=1/4.

Run from any directory: python /path/to/supplement/verify_pell.py
"""

from fractions import Fraction


def require(condition: bool, message: str) -> None:
    if not condition:
        raise AssertionError(message)


def divide_one_minus_z(q: list[int]) -> list[int]:
    partial_sums = []
    total = 0
    for coefficient in q:
        total += coefficient
        partial_sums.append(total)
    require(total == 0, "Division by 1-z left a nonzero remainder")
    return partial_sums[:-1]


def main() -> None:
    u, w = 3, 2
    previous_cluster_gap = 0
    for n in range(2, 8):
        u, w = 3 * u + 4 * w, 2 * u + 3 * w
        require(u * u - 2 * w * w == 1, "Pell equation failed")
        require(u % 2 == 1 and w % 2 == 0, "Required parity failed")
        a, b = (u - 1) // 2, (u - w - 1) // 2
        positive = [0, a, a + 1, u]
        negative = [b, b + 1, u - b - 1, u - b]
        require(len(set(positive + negative)) == 8, "Supports intersect")
        require(min(positive + negative) == 0, "Invalid exponent")
        for r in range(4):
            require(sum(x ** r for x in positive)
                    == sum(x ** r for x in negative),
                    f"n={n}: moment {r} failed")
        fourth = sum(x ** 4 for x in positive) - sum(x ** 4 for x in negative)
        require(4 * fourth == w * w * (w * w - 4) and fourth > 0,
                f"n={n}: exact fourth-moment formula failed")
        clusters = [[0], [b, b + 1], [a, a + 1],
                    [u - b - 1, u - b], [u]]
        cluster_gap = min(clusters[j + 1][0] - clusters[j][-1]
                          for j in range(len(clusters) - 1))
        require(cluster_gap > previous_cluster_gap,
                "Successive computed cluster separations did not increase")
        previous_cluster_gap = cluster_gap
        if n <= 3:
            q = [0] * (u + 1)
            for x in positive:
                q[x] = 1
            for x in negative:
                q[x] = -1
            require(sum(x * x for x in q) == 8, "Wrong coefficient norm")
            p = q
            for _ in range(4):
                p = divide_one_minus_z(p)
            require(p[-1] == 1 and len(p) - 1 == u - 4,
                    "Quotient is not monic of the asserted degree")
            require(24 * sum(p) == fourth, "Exact zero-order check failed")
        print(f"Pell n={n}: u={u}, w={w}, degree={u - 4}, "
              f"moments 0..3=OK, fourth moment={fourth}, "
              f"minimum intercluster gap={cluster_gap}")
    alpha = Fraction(1, 4)
    c1_over_c0 = -alpha / (1 + alpha)
    limit_over_c0 = 8 + 6 * c1_over_c0
    require(limit_over_c0 == Fraction(34, 5), "Limiting rational factor failed")
    print(f"Limiting energy at alpha=1/4, divided by c0(alpha): {limit_over_c0}")
    print("All Pell/PTE integer and rational checks passed.")


if __name__ == "__main__":
    main()