"""Numerical check of the Tunnell-count identities of the paper
"Class Numbers of Imaginary Quadratic Orders in the Polynomial Hierarchy".

For square-free n = 2 (mod 8):  D_n = 2 h(-4n)
For square-free n = 6 (mod 8):  D_n = 0

where D_n = #{(x, y, z) in Z^3 : n = 8x^2 + 2y^2 + 16z^2} and h(-4n) is the
number of reduced primitive positive definite binary quadratic forms of
discriminant -4n.  Exact integer arithmetic only; standard library only
(Python 3.12).

Usage:  python verify_tunnell.py
"""

import random
import time
from math import gcd, isqrt

EXHAUSTIVE_BOUND = 30_000   # all square-free n = 2 (mod 8) up to this bound
RANDOM_BOUND = 200_000      # range of the pseudorandom sample and of the n = 6 (mod 8) check
RANDOM_SAMPLES = 1500
SEED = 2347


def is_squarefree(n):
    """Trial division."""
    p = 2
    while p * p <= n:
        if n % (p * p) == 0:
            return False
        p += 1
    return True


def tunnell_D(n):
    """Number of (x, y, z) in Z^3 with n = 8x^2 + 2y^2 + 16z^2."""
    count = 0
    z = 0
    while 16 * z * z <= n:
        wz = 1 if z == 0 else 2          # z and -z
        rest_z = n - 16 * z * z
        x = 0
        while 8 * x * x <= rest_z:
            wx = 1 if x == 0 else 2      # x and -x
            rest = rest_z - 8 * x * x    # must equal 2y^2
            if rest % 2 == 0:
                q = rest // 2
                y = isqrt(q)
                if y * y == q:
                    wy = 1 if y == 0 else 2
                    count += wz * wx * wy
            x += 1
        z += 1
    return count


def class_number(d):
    """h(-d): number of reduced primitive forms (a, b, c) with b^2 - 4ac = -d."""
    assert d >= 3 and (-d) % 4 in (0, 1)
    count = 0
    a = 1
    while 3 * a * a <= d:                # a <= sqrt(d/3)
        for b in range(-a + 1, a + 1):   # -a < b <= a
            if (b - d) % 2 != 0:         # b = d (mod 2)
                continue
            num = b * b + d
            if num % (4 * a) != 0:
                continue
            c = num // (4 * a)
            if c < a:
                continue
            if a == c and b < 0:
                continue
            if gcd(a, gcd(b, c)) != 1:
                continue
            count += 1
        a += 1
    return count


def main():
    t0 = time.time()

    # 1. exhaustive check of D_n = 2 h(-4n)
    checked = failures = 0
    for n in range(2, EXHAUSTIVE_BOUND + 1, 8):
        if is_squarefree(n):
            checked += 1
            if tunnell_D(n) != 2 * class_number(4 * n):
                failures += 1
                print(f"  FAIL n = {n}: D_n = {tunnell_D(n)}, 2h(-4n) = {2 * class_number(4 * n)}")
    print(f"[1] square-free n = 2 (mod 8), n <= {EXHAUSTIVE_BOUND}: "
          f"{checked} integers checked, {failures} failures of D_n = 2h(-4n)")

    # 2. pseudorandom sample
    rng = random.Random(SEED)
    candidates = [n for n in range(2, RANDOM_BOUND + 1, 8) if is_squarefree(n)]
    sample = rng.sample(candidates, RANDOM_SAMPLES)
    failures2 = 0
    for n in sample:
        if tunnell_D(n) != 2 * class_number(4 * n):
            failures2 += 1
            print(f"  FAIL n = {n}")
    print(f"[2] {RANDOM_SAMPLES} pseudorandom square-free n = 2 (mod 8), n <= {RANDOM_BOUND} "
          f"(seed {SEED}, min {min(sample)}, max {max(sample)}): {failures2} failures of D_n = 2h(-4n)")

    # 3. local vanishing
    checked3 = failures3 = 0
    for n in range(6, RANDOM_BOUND + 1, 8):
        if is_squarefree(n):
            checked3 += 1
            if tunnell_D(n) != 0:
                failures3 += 1
                print(f"  FAIL n = {n}: D_n = {tunnell_D(n)}")
    print(f"[3] square-free n = 6 (mod 8), n <= {RANDOM_BOUND}: "
          f"{checked3} integers checked, {failures3} failures of D_n = 0")

    ok = failures == 0 and failures2 == 0 and failures3 == 0
    print("ALL CHECKS PASSED" if ok else "SOME CHECKS FAILED")
    print(f"elapsed: {time.time() - t0:.1f} s")


if __name__ == "__main__":
    main()