"""Independent numerical and exact checks for numberdb.org/T643.

Download generate.py, checks.py, check.py and companions.json beside each other.
Run: sage -python checks.py
Set NUMBERDB_CHECK_STORED=1 (and credentials for the draft) to read every value
back from NumberDB before repeating all checks. The campaign invokes main with
stored=True directly. All evidence is printed as JSON; redirect stdout to save it.

The independent exact path uses Zagier section 2's B(d) recurrences and section
6's half-integral-weight Hecke operators, without class polynomials or Faber
polynomials. A second path enumerates reduced integral forms and sums Hecke
translates of j in 1024-bit Arb arithmetic. Known-answer and wrong-answer
controls precede comparisons. Companion documents are read-only snapshots.
"""
import json
import os
import sys
from functools import lru_cache
from math import gcd, isqrt
from pathlib import Path

import numberdb.sage as numberdb
from sage.rings.integer_ring import ZZ
from sage.rings.rational_field import QQ
from sage.rings.complex_arb import ComplexBallField
from sage.rings.polynomial.polynomial_ring_constructor import PolynomialRing
from sage.version import version as sage_version
from sage.arith.misc import kronecker_symbol
from check import exactness, measure, agrees_with, names_its_rings
import generate

PRECISION_BITS = 1024
B = {-1: ZZ(1), 0: ZZ(-2)}


def zagier_coefficients(bound):
    for n in range(1, (bound + 3) // 4 + 1):
        B[4*n-1] = 240 * sum(ZZ(a)**3 for a in range(1, n+1) if n % a == 0) - sum(
            r*r*B.get(4*n-r*r, ZZ(0)) for r in range(2, isqrt(4*n+1)+1))
        B[4*n] = -2 * sum(B.get(4*n-r*r, ZZ(0)) for r in range(1, isqrt(4*n+1)+1))


def coefficient(d):
    if d > max(B):
        raise ValueError('Insufficient independent coefficient bound')
    return B.get(d, ZZ(0))


def prime_hecke(d, p):
    if d < 3 or d % 4 not in (0, 3):
        return ZZ(0)
    return (coefficient(p*p*d) + kronecker_symbol(-d, p)*coefficient(d)
            + (p*coefficient(d//(p*p)) if d % (p*p) == 0 else 0))


def independent_trace(key):
    m, d = key
    if m == 1:
        return -coefficient(d)
    if m in (2, 3):
        return -prime_hecke(d, m)
    if m == 4:
        return -(prime_hecke(4*d, 2) + kronecker_symbol(-d, 2)*prime_hecke(d, 2)
                 + (2*prime_hecke(d//4, 2) if d % 4 == 0 else 0) - 2*coefficient(d))
    raise ValueError('Independent check implemented for orders 1 through 4')


def reduced_forms(d):
    for a in range(1, isqrt(d//3)+1):
        for b in range(-a, a+1):
            if (b*b+d) % (4*a):
                continue
            c = (b*b+d)//(4*a)
            if a > c or ((abs(b) == a or a == c) and b < 0):
                continue
            yield a, b, c


def weight(a, b, c):
    return 3 if a == b == c else 2 if b == 0 and a == c else 1


@lru_cache(None)
def hecke_j_sum(m, a, b, c):
    C = ComplexBallField(PRECISION_BITS)
    d = 4*a*c - b*b
    tau = (C(-b) + C(0, 1)*C(d).sqrt())/(2*a)
    total = C(0)
    for denominator in range(1, m+1):
        if m % denominator:
            continue
        numerator = m//denominator
        for shift in range(denominator):
            total += ((numerator*tau+shift)/denominator).modular_j() - 744
    return total


def numerical_trace(m, d):
    C = ComplexBallField(PRECISION_BITS)
    total = C(0)
    for a, b, c in reduced_forms(d):
        total += hecke_j_sum(m, a, b, c)/weight(a, b, c)
    return total


def scalar(record):
    return record['number'] if isinstance(record, dict) else record


def stored_values(document):
    result = {}
    entries = document['Numbers']
    if isinstance(entries, list):
        for row in entries:
            result[(int(row['params']['m']), int(row['params']['d']))] = ZZ(row['number'])
    else:
        for m, block in entries.items():
            for d, row in block.items():
                result[(int(m), int(d))] = ZZ(scalar(row))
    return result


def g_product(bound):
    """Coefficients of q*g: theta_1 * E4(q^4) * product(1-q^(4n))^-6."""
    v = [ZZ(1)] + [ZZ(0)]*bound
    for n in range(4, bound+1, 4):
        for repetition in range(6):
            for k in range(n, bound+1):
                v[k] += v[k-n]
    e = [ZZ(1)] + [ZZ(0)]*bound
    for a in range(1, bound//4+1):
        for k in range(4*a, bound+1, 4*a):
            e[k] += 240*ZZ(a)**3
    product = [sum(v[k-j]*e[j] for j in range(k+1)) for k in range(bound+1)]
    return [product[k] + sum(2*(-1)**r*product[k-r*r]
            for r in range(1, isqrt(k)+1)) for k in range(bound+1)]


def main(stored=False):
    generator = generate.ZagierTraces()
    expected_keys = [(m,d) for m in range(1,5) for d in range(3,201) if d%4 in (0,3)]
    assert [(p['m'],p['d']) for p in generator.enumerate()] == expected_keys
    if stored:
        document = numberdb.table('T643')
        values = stored_values(document)
    else:
        values = {key:generate.trace(*key) for key in expected_keys}
    assert list(values) == expected_keys
    assert not exactness(values)
    assert not names_its_rings(str(Path(__file__).with_name('generate.py')))
    zagier_coefficients(3200)
    printed = {3:-248,4:492,7:-4119,8:7256,11:-33512,12:53008,
               15:-192513,16:287244,19:-885480,20:1262512,23:-3493982,
               24:4833456,27:-12288992,28:16576512,31:-39493539,32:52255768}
    for d, answer in printed.items():
        assert independent_trace((1,d)) == answer
        assert values[1,d] == answer
    assert agrees_with({(1,3):ZZ(-247)}, independent_trace), 'Wrong-answer control failed'
    disagreements = agrees_with(values, independent_trace)
    assert not disagreements, disagreements[:10]
    product = g_product(201)
    assert product[0] == 1 and product[1] == -2
    for d in range(3,201):
        assert product[d+1] == coefficient(d)
    C = ComplexBallField(PRECISION_BITS)
    controls = [(C(0,1).modular_j(),ZZ(1728)),
                (((-1+C(0,1)*C(3).sqrt())/2).modular_j(),ZZ(0))]
    for ball, answer in controls:
        assert (ball.real().is_finite() and ball.imag().is_finite()) and (answer in ball)
        assert ball.real().rad() < QQ(1)/10**200 and ball.imag().rad() < QQ(1)/10**200
        assert not (answer+1 in ball)
    max_real_radius = 0.0
    max_imag_radius = 0.0
    for (m,d), value in values.items():
        ball = numerical_trace(m,d)
        assert (ball.real().is_finite() and ball.imag().is_finite()) and (value in ball), (m,d,str(ball),str(value))
        assert ball.real().rad() < QQ(1)/10**110
        assert ball.imag().rad() < QQ(1)/10**110
        assert not (value+1 in ball), (m,d,'wrong answer accepted')
        max_real_radius = max(max_real_radius, float(ball.real().rad()))
        max_imag_radius = max(max_imag_radius, float(ball.imag().rad()))
    companions = json.loads(Path(__file__).with_name('companions.json').read_text())
    Zx = PolynomialRing(ZZ, 'x'); x = Zx.gen()
    ZX = PolynomialRing(ZZ, 'X')
    for d in range(3,201):
        if d%4 not in (0,3):
            continue
        H = generate.class_polynomial(d)
        assert H == Zx(scalar(companions['T129']['Numbers'][str(-d)]))
        forms = list(reduced_forms(d))
        assert H.degree() == sum(gcd(gcd(a,b),c) == 1 for a,b,c in forms)
        # H(d) and the j versus J conversion, using explicit CM j balls.
        weighted_class_number = sum(QQ(1)/weight(*form) for form in forms)
        raw_j_trace = sum((hecke_j_sum(1,*form)+744)/weight(*form) for form in forms)
        assert (values[1,d] + 744*weighted_class_number in raw_j_trace)
        # d=12 and d=16 specifically exercise imprimitive elliptic weights.
    assert list(reduced_forms(12)) == [(1,0,3),(2,2,2)]
    assert list(reduced_forms(16)) == [(1,0,4),(2,0,2)]
    for m in range(2,5):
        assert generate.faber(m) == ZX(scalar(companions['T635']['Numbers']['1A'][str(m)]))
    printed_polys = {2:x*x-1488*x+159768,
                     3:x**3-2232*x*x+1069956*x-36866976,
                     4:x**4-2976*x**3+2533680*x*x-561444608*x+8507424792}
    for m,F in printed_polys.items():
        assert Zx(generate.faber(m)(x-744)) == F
    for n,c in enumerate(generate.j_positive_coefficients(3),1):
        assert c == ZZ(scalar(companions['T79']['Numbers'][str(n)]))
    for D,row in companions['T83']['Numbers'].items():
        d=-int(D); j=ZZ(scalar(row))
        assert generate.class_polynomial(d) == x-j
        for m in range(1,5):
            assert generate.primitive_trace(m,d) == generate.faber(m)(j-744)
    # Full principal parts q^-m+O(q), checked against the independently stored j coefficients.
    from sage.rings.laurent_series_ring import LaurentSeriesRing
    L=LaurentSeriesRing(QQ,'q'); q=L.gen()
    J=q**(-1)+sum(ZZ(scalar(companions['T79']['Numbers'][str(k)]))*q**k for k in range(1,9))
    for m in range(1,5):
        series=generate.faber(m)(J)
        assert all(series[e] == (1 if e == -m else 0) for e in range(-m,1))
    report = {'table':'T643','stored_readback':stored,'entries':len(values),
              'exact_hecke_comparisons':len(values),'direct_CM_ball_comparisons':len(values),
              'precision_bits':PRECISION_BITS,'max_real_radius_upper_diagnostic':max_real_radius,
              'max_imag_radius_upper_diagnostic':max_imag_radius,
              'rigorous_absolute_radius_bound':'< 1e-110 in each component',
              'source_printed_trace_controls':len(printed),'explicit_source_polynomials':3,
              'class_polynomials_vs_T129':100,'rational_CM_controls_T83':13,
              'faber_polynomials_vs_T635':3,'g_product_coefficients_checked_through':200,
              'known_and_wrong_answer_controls':'passed', 'measure':measure(values),
              'sage_version':sage_version,
              'skill_manifest':'agents/runs/20261101T063157Z-build-skills/manifest.json',
              'skill_sha256':'a2b28d9953264d7b0ca62663785c2952190c5be6d6d3597fd695d9a54d606278',
              'values':{str(m)+','+str(d):str(v) for (m,d),v in values.items()}}
    print(json.dumps(report,indent=2,default=str))
    return report


if __name__ == '__main__':
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN') == '1':
        numberdb.configure(api_key=sys.stdin.read().strip())
    main(stored=os.environ.get('NUMBERDB_CHECK_STORED') == '1')
