"""Independent checks for T529. Run: sage -python checks.py.

Keep generate.py and check.py beside this file. With NUMBERDB_API_KEY set,
use --stored to check database readback too. NUMBERDB_KEY_FROM_STDIN=1 reads
the key from stdin and also selects readback, for agents/sage.sh.
"""
import json
import os
import sys
import hashlib
import platform
from pathlib import Path
from decimal import Decimal
from fractions import Fraction
from math import factorial
import numberdb.sage as numberdb
from sage.rings.rational_field import QQ
from sage.rings.real_arb import RealBallField
from sage.rings.polynomial.polynomial_ring_constructor import PolynomialRing
from numberdb._write import to_text
from sage.version import version as sage_version
from generate import PositiveBridgeMoments, coefficients, rational_factor, WORKING_GUARD
from check import exactness, measure, agrees_with, names_its_rings

LIMIT = 12
HIGH_BITS = 768
R = RealBallField(HIGH_BITS)


def resolvent_coefficients(limit):
    """Direct polynomial solution of the half-line Green-function problem."""
    ring = PolynomialRing(QQ, 'x')
    x = ring.gen()
    polynomials = [ring(1)]
    origins = [QQ(1)]
    for n in range(1, limit+1):
        rhs = 2*x*polynomials[-1]
        degree = int(rhs.degree())
        # For q=p': 2q-q'=rhs, solved from highest coefficient down.
        q = [QQ(0)]*(degree+2)
        for d in range(degree, -1, -1):
            q[d] = (rhs[d]+QQ(d+1)*q[d+1])/QQ(2)
        p = ring([q[0]/QQ(2)] + [q[d]/QQ(d+1) for d in range(degree+1)])
        assert 2*p.derivative()-p.derivative(2) == rhs
        assert p.derivative()(0) == 2*p(0)
        polynomials.append(p)
        origins.append(p(0))
    return origins, polynomials


def airy_quotient(limit):
    """Formal a(z)/h(z) in Janson (242), from factorial coefficients."""
    a = [QQ(factorial(6*n))/QQ(576**n*factorial(3*n)*factorial(2*n))
         for n in range(limit+1)]
    h = [a[n]/QQ(1-6*n) for n in range(limit+1)]
    answer = []
    for n in range(limit+1):
        answer.append(a[n]-sum((h[j]*answer[n-j] for j in range(1,n+1)),QQ(0)))
    return answer


def direct_moment(n, origin):
    return R(origin)*R.pi().sqrt()*R(factorial(n))/(R(2)**(R(3*n)/R(2))
            *(R(3*n+1)/R(2)).gamma())


def rational_decimal(text):
    f = Fraction(Decimal(text))
    return QQ(f.numerator)/QQ(f.denominator)


def bounds(text):
    if '.' not in text and 'e' not in text.lower():
        q = QQ(text)
        return q, q
    decimal = Decimal(text)
    center = rational_decimal(text)
    exponent = decimal.as_tuple().exponent
    unit = QQ(10)**int(exponent)
    return center-unit, center+unit


def interval(text):
    lo, hi = bounds(text)
    return R((lo+hi)/QQ(2), (hi-lo)/QQ(2))


def assert_contains(text, value):
    assert value.is_finite()
    assert value.accuracy() >= 700
    lo, hi = bounds(text)
    if lo == hi:
        assert lo in value
    else:
        assert value > R(lo) and value < R(hi)


def run(stored_document=None):
    origins0, polys0 = resolvent_coefficients(1)
    assert origins0 == [QQ(1), QQ(1)/4]
    assert polys0[1] == polys0[0].parent()([QQ(1)/4, QQ(1)/2, QQ(1)/2])
    assert 1 in direct_moment(0,origins0[0])
    # The marginal Gaussian mean integrates to Beta(3/2,3/2)/sqrt(2*pi).
    gaussian_mean = (R(3)/2).gamma()**2/R(3).gamma()/(2*R.pi()).sqrt()
    assert gaussian_mean.is_finite() and gaussian_mean.accuracy() >= 700
    assert 0 in (gaussian_mean-direct_moment(1,origins0[1]))
    generator = PositiveBridgeMoments()
    path = Path(__file__).with_name('generate.py')
    assert names_its_rings(str(path)) == []
    enumeration = list(generator.enumerate())
    assert enumeration == [{'n':str(n)} for n in range(1,LIMIT+1)]
    values = {int(p['n']):generator.value(p,100) for p in enumeration}
    assert exactness(values) == []
    origins, polynomials = resolvent_coefficients(LIMIT)
    k, p = coefficients(LIMIT)
    independent_p = {n:origins[n]/QQ(2**n) for n in range(LIMIT+1)}
    assert agrees_with(dict(enumerate(p)),lambda n:independent_p[n]) == []
    assert agrees_with(dict(enumerate(p)),lambda n:airy_quotient(LIMIT)[n]) == []
    # Printed exact values of Table 6, transcribed before running this check.
    table6 = {0:QQ(1),1:QQ(1)/8,2:QQ(1)/20,3:QQ(71)/4096,4:QQ(211)/18480}
    assert agrees_with({n:rational_factor(n) for n in range(5)},lambda n:table6[n]) == []
    assert table6[2] == (QQ(7)/60+QQ(1)/12)/4
    # These checks are on exact rational coefficients, not ball overlap.
    assert exactness(dict(enumerate(k))) == []
    assert exactness(dict(enumerate(p))) == []
    texts = {n:to_text(v,digits=100) for n,v in values.items()}
    rows = []
    minimum_bits = min(int(values[n].accuracy()) for n in values if n%2)
    assert minimum_bits >= 500
    for n,text in texts.items():
        q = rational_factor(n)
        independent = direct_moment(n,origins[n])
        assert_contains(text,independent)
        if n%2:
            assert len(Decimal(text).as_tuple().digits) == 100
            assert values[n].is_finite()
            lo,hi=bounds(text)
            assert values[n] > lo and values[n] < hi
        else:
            assert QQ(text) == q
        rows.append({'n':n,'rational_factor':str(q),'factor':'sqrt(pi/2)' if n%2 else '1',
                     'P_n':str(p[n]),'resolvent_origin':str(origins[n]),'number':text,
                     'working_accuracy_bits':int(values[n].accuracy()) if n%2 else None,
                     'independent_accuracy_bits':int(independent.accuracy())})
    check_written_identities(texts, k, p)
    if stored_document is not None:
        block = stored_document['Numbers']
        if isinstance(block,dict):
            saved={int(n):v['number'] if isinstance(v,dict) else v for n,v in block.items()}
        else:
            saved={int(v['params']['n']):v['number'] for v in block}
        assert list(saved) == list(range(1,LIMIT+1))
        assert saved == texts
        for n,text in saved.items(): assert_contains(text,direct_moment(n,origins[n]))
        check_written_identities(saved,k,p)
        # Test the public one-value program independently of the generator interface.
        namespace={}
        exec(stored_document['Programs']['raw-moment']['code'],namespace)
        for n in range(1,LIMIT+1):
            v=namespace['moment'](n)
            if n%2:
                assert v.is_finite() and v.accuracy() >= 500
                lo,hi=bounds(saved[n]); assert v>lo and v<hi
            else: assert v == QQ(saved[n])
    manifest_path=Path(__file__).with_name('skill-manifest.json')
    manifest=json.loads(manifest_path.read_text())
    return {'table':'T529','status':'pass','stored_readback_checked':stored_document is not None,
            'skill_manifest':'agents/runs/20261026T213949Z-build-skills/manifest.json',
            'skill_manifest_sha256':hashlib.sha256(manifest_path.read_bytes()).hexdigest(),
            'skill_versions':manifest['versions'],
            'versions':{'python':platform.python_version(),'sage':sage_version,
                        'numberdb':getattr(numberdb,'__version__','unknown')},
            'selection':{'parameter':'n','type':'Z','domain':'n >= 1','enumeration':list(range(1,13)),
                         'category':'discrete_index','independent_critique_pending':True},
            'working_bits':numberdb.bits(100,losing=WORKING_GUARD),'guard_bits':WORKING_GUARD,
            'independent_bits':HIGH_BITS,'minimum_working_accuracy_bits':minimum_bits,
            'written_measurement':measure(texts),'ball_measurement':measure(values),
            'checks':{'exactness':True,'named_imports':True,'resolvent_polynomial_identities':13,
                      'resolvent_controls':'M0=1; p1=x^2/2+x/2+1/4',
                      'exact_resolvent_coefficient_comparisons':13,'exact_airy_quotient_comparisons':13,
                      'table6_moment_comparisons':5,'gaussian_marginal_mean':True,
                      'second_moment_from_signed_and_absolute_bridge':True,
                      'decimal_interval_containment':6,'exact_rational_rows':6,
                      'written_recurrence_checks':12,'written_strict_log_convexity_checks':11},
            'rows':rows}


def check_written_identities(texts,k,p):
    moments = {0:R(1),**{n:interval(text) for n,text in texts.items()}}
    recovered={0:R(1)}
    for n in texts:
        recovered[n]=moments[n]*(R(3*n+1)/2).gamma()*R(2)**(R(n)/2)/(R.pi().sqrt()*R(factorial(n)))
        assert recovered[n].is_finite() and recovered[n].accuracy() >= 310
        assert p[n] in recovered[n]
        rhs=sum((R(k[j])*recovered[n-j] for j in range(1,n+1)),R(0))
        residual=recovered[n]-rhs
        assert residual.is_finite() and 0 in residual
        assert residual.rad_as_ball() < abs(R(p[n]))*R(2)**(-310)
    for n in range(1,LIMIT):
        assert moments[n-1]*moments[n+1] > moments[n]**2


if __name__ == '__main__':
    from_stdin=os.environ.get('NUMBERDB_KEY_FROM_STDIN') == '1'
    if from_stdin: numberdb.configure(api_key=sys.stdin.read().strip())
    doc=numberdb.table('T529') if from_stdin or '--stored' in sys.argv else None
    print(json.dumps(run(doc),indent=2))
