"""Exact finite checks for the BF author-edition singular-series counterexample.

Standard library only; no manuscript code is executed. Run with python3 and
redirect stdout to JSON. Finite checks support, but do not replace, the proof
in singular-series-counterexample.md. Counts use the intended p^k domain.
"""
from collections import Counter
from fractions import Fraction
from math import gcd
import json
import platform


def distribution(coefficients, modulus):
    a, b, c = coefficients
    result = [0] * modulus
    for x in range(modulus):
        for y in range(modulus):
            result[(a*x*x+b*x*y+c*y*y) % modulus] += 1
    assert sum(result) == modulus**2
    return result


def form(n):
    return (n*n+1, -2*n, n*n)


def count(t, p, k):
    n = 5**t
    q = p**k
    first, second = distribution(form(n), q), distribution(form(2*n), q)
    m = -3*n*n
    total = sum(first[s] * second[(s-m) % q] for s in range(q))
    return {"t": t, "p": p, "k": k, "modulus": q,
            "solutions": total, "normalized": str(Fraction(total, q**3))}


def product_count_checks():
    checks = 0
    for p in (2, 3, 5, 7):
        for k in range(1, 4):
            q = p**k
            dist = Counter((z*w) % q for z in range(q) for w in range(q))
            for r in range(k):
                b = p**r
                assert dist[b] == (r+1)*(q-q//p)
                checks += 1
    return checks


def main():
    rows = []
    for t in (1, 2, 3):
        n = 5**t
        a, ap = n*n, 4*n*n
        assert gcd(gcd(*form(n)[:2]), form(n)[2]) == 1
        assert form(n)[0]*form(n)[2] - (form(n)[1]//2)**2 == a*a
        assert (n-1)**2+2*(n-1)+2-1 == a
        assert gcd(n-1, 1) == 1
        assert 2*(0+1+(n+1)**2)-n*n == (n+2)**2
        # Known prime support of aa'(a-a') is {2,3,5}; common support {5}.
        assert Fraction(2)*Fraction(3, 2)*Fraction(4, 3) == 4
        for p, ks in ((2, range(3, 8)), (3, range(2, 5)),
                      (7, range(1, 3)), (11, range(1, 3))):
            for k in ks:
                row = count(t, p, k)
                expected = Fraction(2) if p == 2 else (
                    Fraction(32, 27) if p == 3 else 1-Fraction(1, p*p))
                assert Fraction(row['normalized']) == expected, row
                rows.append(row)
    # Observe the transient increase, then the predicted stable 5-adic count.
    for t, ks in ((1, range(1, 5)), (2, range(1, 6))):
        for k in ks:
            row = count(t, 5, k)
            if k > 2*t:
                assert Fraction(row['normalized']) == Fraction(4*(2*t+1), 5)
            rows.append(row)
    assert [r['solutions'] for r in rows if r['t']==1 and r['p']==5][:3] == [225, 40625, 4687500]
    # Entire Euler product coefficient in S=(coefficient/pi^2)(2t+1).
    assert Fraction(2)*Fraction(32,27)*Fraction(4,5)*Fraction(75,8) == Fraction(160,9)
    print(json.dumps({"arithmetic": "exact Python integers and Fraction",
                      "python": platform.python_version(),
                      "method": "binary residue histograms joined by F=m",
                      "product_count_cases": product_count_checks(),
                      "local_count_cases": len(rows), "rows": rows,
                      "all_assertions_passed": True}, indent=2))


if __name__ == '__main__':
    main()