"""Independent checks for Page average entropies -- numberdb.org/T573.

Run: sage -python verify_math.py
Read back the draft too: sage -python verify_math.py --stored
Requires generate.py and check.py beside this file, plus Sage and numberdb.
NUMBERDB_CHECK_STORED=1 also enables readback. Authenticate as for generate.py.

The entropy check expands the monic generalized Laguerre polynomials in
Sen (1996), equations (2)-(3), and integrates their squared one-point kernel.
No Page harmonic-difference formula or generator helper is used to obtain the
check values. Factorial moments and their logarithmic derivatives are exact.
Every kernel is first checked for mass N and first moment NK. The two-qubit
control separately integrates 3(2*x-1)^2 on [0,1], giving entropy 1/3.

For bits, a positive atanh series and a geometric tail bound enclose log(2)
in rational arithmetic independently of Arb's logarithm implementation.
Stored decimals mean center +/- one unit in the last written decimal place.
"""
import json
import os
import sys
from decimal import Decimal
from fractions import Fraction
from functools import lru_cache
from math import comb, factorial
from pathlib import Path

import numberdb.sage as numberdb
from numberdb import __version__ as NUMBERDB_VERSION
from numberdb._write import to_text
from sage.rings.rational_field import QQ
from sage.rings.real_mpfi import RealIntervalField
from sage.env import SAGE_VERSION

from generate import PageEntropies, WORKING_GUARD
from check import exactness, measure, agrees_with, names_its_rings


@lru_cache(None)
def harmonic(j):
    return sum((Fraction(1, i) for i in range(1, j + 1)), Fraction(0))


@lru_cache(None)
def laguerre_entropy(N, K):
    alpha = K - N
    mass, trace, logmoment = Fraction(0), Fraction(0), Fraction(0)
    for degree in range(N):
        coefficients = [Fraction((-1)**(degree - i) * comb(degree, i) * factorial(degree + alpha),
                                 factorial(alpha + i)) for i in range(degree + 1)]
        norm = factorial(degree) * factorial(degree + alpha)
        polynomial_mass, polynomial_trace, logarithmic_part = Fraction(0), Fraction(0), Fraction(0)
        for i, a in enumerate(coefficients):
            for j, b in enumerate(coefficients):
                r = alpha + i + j
                coefficient = a * b / norm
                polynomial_mass += coefficient * factorial(r)
                polynomial_trace += coefficient * factorial(r + 1)
                logarithmic_part += coefficient * factorial(r + 1) * harmonic(r + 1)
        assert polynomial_mass == 1
        assert polynomial_trace == 2 * degree + alpha + 1
        mass += polynomial_mass
        trace += polynomial_trace
        logmoment += logarithmic_part
    assert mass == N and trace == N * K
    # The coefficient of Euler's constant cancels because trace == N*K.
    return harmonic(N * K) - logmoment / (N * K)


def controls():
    # Integral of x^r and of -x^r log(x) on [0,1].
    coeffs = (3, -12, 12)
    density_mass = sum((Fraction(c, r + 1) for r, c in enumerate(coeffs)), Fraction(0))
    entropy = 2 * sum((Fraction(c, (r + 2)**2) for r, c in enumerate(coeffs)), Fraction(0))
    assert density_mass == 1 and entropy == Fraction(1, 3)
    assert laguerre_entropy(2, 2) == entropy
    assert laguerre_entropy(2, 3) == Fraction(9, 20)
    return {'two_qubit_density_mass': str(density_mass), 'two_qubit_entropy_nats':str(entropy),
            'two_by_three_entropy_nats': '9/20'}


def log2_bounds(terms=150):
    t = Fraction(1, 3)
    low = 2 * sum((t**(2*j + 1) / (2*j + 1) for j in range(terms)), Fraction(0))
    tail = 2 * t**(2*terms + 1) / ((2*terms + 1) * (1 - t*t))
    return low, low + tail


def exact_fraction(value):
    return Fraction(str(value))


def ball_bounds(ball):
    assert ball.is_finite()
    interval = RealIntervalField(ball.parent().precision() + 16)(ball)
    return (exact_fraction(interval.lower().exact_rational()),
            exact_fraction(interval.upper().exact_rational()))


def decimal_bounds(text):
    center = Fraction(text)
    ulp = Fraction(10)**Decimal(text).as_tuple().exponent
    return center - ulp, center + ulp


def default_selection():
    # Independent exhaustive interpretation of the stated dimension budget.
    return [{'N':N, 'K':K, 'unit':unit}
            for N in range(2, 129) for K in range(N, 129)
            if N*K <= 128 for unit in ('nats', 'bits')]


def check_records(records, loglo, loghi):
    expected_params = default_selection()
    assert len(records) == len(expected_params) == 400
    for record, params in zip(records, expected_params):
        assert {k:str(v) for k,v in record['params'].items()} == {k:str(v) for k,v in params.items()}
        exact = laguerre_entropy(params['N'], params['K'])
        text = str(record['number'])
        if params['unit'] == 'nats':
            assert '.' not in text and 'e' not in text.lower()
            assert Fraction(text) == exact
        else:
            lower, upper = decimal_bounds(text)
            assert lower <= exact/loghi <= exact/loglo <= upper
            assert 0 < (upper-lower)/lower < Fraction(1, 10**98)
    return len(records)


def check_program(document):
    namespace={}
    exec(document['Programs']['program-sage']['code'],namespace)
    loglo,loghi=log2_bounds()
    records=[]
    for params in default_selection():
        value=namespace['page_entropy'](params['N'],params['K'],params['unit'])
        records.append({'params':params,'number':to_text(value,digits=100)})
    return check_records(records,loglo,loghi)


def run_checks(stored=False):
    report = {'controls': controls()}
    generator = PageEntropies()
    parameters = list(generator.enumerate())
    assert parameters == default_selection()
    assert not names_its_rings(Path(__file__).with_name('generate.py'))
    values = {(p['N'],p['K'],p['unit']):generator.value(p,100) for p in parameters}
    assert not exactness(values)
    nats = {key:value for key,value in values.items() if key[-1]=='nats'}
    assert not agrees_with(nats, lambda key: QQ(laguerre_entropy(key[0],key[1])))
    loglo, loghi = log2_bounds()
    assert 0 < loghi-loglo < Fraction(1,10**145)
    # Sanity control for a positive, nonzero logarithm interval.
    assert Fraction(69,100) < loglo < loghi < Fraction(70,100)
    smallest_accuracy = None
    records = []
    monotonic_pairs = 0
    for p in parameters:
        N,K,unit = p['N'],p['K'],p['unit']
        value = values[(N,K,unit)]
        text = to_text(value,digits=100)
        records.append({'params':p, 'number':text})
        if unit=='bits':
            lower,upper = ball_bounds(value)
            assert 0 < lower < upper
            exact = laguerre_entropy(N,K)
            assert lower <= exact/loghi <= exact/loglo <= upper
            assert (upper-lower)/lower < Fraction(1,10**117)
            stored_lo,stored_hi=decimal_bounds(text)
            assert stored_lo <= lower <= upper <= stored_hi
            relative_width = (upper-lower)/lower
            accuracy = relative_width.denominator.bit_length() - relative_width.numerator.bit_length()
            if relative_width >= Fraction(1, 2**accuracy):
                accuracy -= 1
            smallest_accuracy=accuracy if smallest_accuracy is None else min(smallest_accuracy,accuracy)
        else:
            field=RealIntervalField(512)
            assert field(value) > 0 and field(value) < field(N).log()
            if (N,K+1,'nats') in values:
                assert value < values[(N,K+1,'nats')]
                monotonic_pairs+=1
    check_records(records, loglo, loghi)
    measured = measure(values)
    report.update(entries=400,dimension_pairs=200,exact_nats=200,certified_bits=200,
                  independent_laguerre_checks=200,kernel_normalization_checks=200,
                  monotonic_K_pairs=monotonic_pairs,entropy_bound_checks=200,
                  independent_log2_series_terms=150,independent_log2_width_less_than='1e-145',
                  minimum_certified_relative_width_bits=smallest_accuracy,
                  arb_working_bits=numberdb.bits(100,losing=WORKING_GUARD),guard_bits=WORKING_GUARD,
                  all_arb_relative_widths_less_than='1e-117',
                  serialization_contains_entire_arb_ball=200,
                  serialization_contains_independent_interval=200,
                  measurement=measured,
                  longest_serialized_value=max(len(r['number']) for r in records),
                  serialized_records_json_bytes=len(json.dumps(records).encode()),
                  versions={'sage':SAGE_VERSION,'numberdb':NUMBERDB_VERSION,'python':sys.version.split()[0]},
                  skill_manifest='agents/runs/20261028T054431Z-build-skills/manifest.json',
                  skill_sha256='e37efb7e2ededf6c1d7134c673f152f648f26e0267c556aebdf3ec65765d8b4c')
    if stored:
        document = numberdb.table('T573')
        report['stored_entries_independently_checked']=check_records(document['Numbers'],loglo,loghi)
        report['stored_program_entries_checked']=check_program(document)
        verify=generator.verify(sample=None)
        assert verify.ok, str(verify)
        report['generator_verify']=str(verify)
    return report


if __name__=='__main__':
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN')=='1':
        os.environ['NUMBERDB_API_KEY']=sys.stdin.read().strip()
    report=run_checks('--stored' in sys.argv or os.environ.get('NUMBERDB_CHECK_STORED')=='1')
    print(json.dumps(report,indent=2))