"""Independent Scorer checks using certified Laplace quadrature and mpmath.

Run as python checks.py --kind gi (or hi). --output-dir writes a numerical
working copy and evidence; it never sends a table. --live compares every
stored entry, using NUMBERDB_API_KEY when the table is still a draft.

Hi is integrated on [0,12]. For T=12 and x<T*T its omitted positive tail
is at most exp(x*T-T**3/3)/(pi*(T*T-x)), because the logarithmic derivative
of the integrand is <= x-T*T on [T,infinity). Gi is checked as Bi-Hi,
independently of the generator's inhomogeneous Taylor series.
"""

import argparse
import hashlib
import importlib.metadata
import json
from decimal import Decimal
from fractions import Fraction
from pathlib import Path

import mpmath as mp
import numberdb
from flint import acb, arb
from numberdb._write import to_text

from scorer import (ScorerValues, argument_for, as_interval,
                    evaluate, mp_fraction, precision, working_precision, widen)


def decimal_interval(text):
    decimal = Decimal(text)
    midpoint = Fraction(decimal)
    unit = Fraction(10) ** decimal.as_tuple().exponent
    return numberdb.RealInterval(midpoint - unit, midpoint + unit)


def contains(outer, inner):
    return outer.lower <= inner.lower and inner.upper <= outer.upper


def quadrature_hi(argument, digits=140):
    tolerance = arb(10) ** (-digits - 10)
    integral = acb(0)
    for left in range(0, 12, 3):
        integral += acb.integral(
            lambda position, analytic: (argument * position - position ** 3 / 3).exp(),
            left, left + 3, rel_tol=tolerance, abs_tol=tolerance)
    integral /= arb.pi()
    tail = (12 * argument - arb(576)).exp() / (arb.pi() * (144 - argument))
    if not argument < 144 or not integral.imag.contains(0):
        raise ArithmeticError('quadrature assumptions failed')
    result = widen(integral.real, tail)
    if not result.is_finite() or result.rad() > arb(10) ** (-digits):
        raise ArithmeticError('quadrature did not reach its error target')
    return result


def quadrature_controls():
    tolerance = arb('1e-150')
    polynomial = acb.integral(lambda position, analytic: position ** 2,
                             0, 3, abs_tol=tolerance, rel_tol=tolerance)
    exponential = acb.integral(lambda position, analytic: (-position).exp(),
                              0, 1, abs_tol=tolerance, rel_tol=tolerance)
    if not polynomial.real.contains(9) or not exponential.real.overlaps(1 - arb(-1).exp()):
        raise ArithmeticError('nonzero quadrature control failed')
    if polynomial.real.rad() > tolerance or exponential.real.rad() > tolerance:
        raise ArithmeticError('quadrature controls too wide')


def check(kind, live=False):
    generator = ScorerValues()
    generator.kind = kind
    generator.table = 'T514' if kind == 'gi' else 'T515'
    stored = numberdb.table(generator.table)['Numbers'] if live else None
    numbers = {}
    records = []
    with precision(working_precision(130)):
        quadrature_controls()
        for params in generator.enumerate():
            identity = params['x']
            result = generator.value(params, 100)
            generated = to_text(result['number'], 100)
            written = generated
            if stored is not None:
                entry = stored[identity]
                written = entry['number'] if isinstance(entry, dict) else entry
            written_interval = decimal_interval(written)
            argument = argument_for(identity, 130)
            integral = quadrature_hi(argument)
            independent = integral if kind == 'hi' else argument.airy_bi() - integral
            if not contains(written_interval, as_interval(independent)):
                raise ArithmeticError('independent enclosure is outside stored digits: ' + identity)
            if not contains(written_interval, result['number']):
                raise ArithmeticError('generator enclosure is outside stored digits: ' + identity)
            function = evaluate(argument, kind, digits=130)
            second = evaluate(argument, kind, derivative=2, digits=130)
            forcing = (-1 if kind == 'gi' else 1) / arb.pi()
            if not (second - argument * function - forcing).contains(0):
                raise ArithmeticError('differential equation check failed: ' + identity)
            if identity.startswith('ai-zero-') and not argument.airy_ai().contains(0):
                raise ArithmeticError('Airy zero residual excludes zero')
            if identity.startswith('gi-critical-') and not evaluate(
                    argument, derivative=1, digits=130).contains(0):
                raise ArithmeticError('critical-point residual excludes zero')
            for decimal_digits in (150, 200):
                with mp.workdps(decimal_digits):
                    if identity.startswith('ai-zero-'):
                        candidate = mp.airyaizero(int(identity[len('ai-zero-'):]))
                    elif identity.startswith('gi-critical-'):
                        midpoint = mp_fraction(as_interval(argument).midpoint())
                        candidate = mp.findroot(lambda position: mp.diff(mp.scorergi, position),
                                                (midpoint - mp.mpf('.001'), midpoint + mp.mpf('.001')))
                    else:
                        candidate = mp.mpf(identity)
                    reference = (mp.scorergi if kind == 'gi' else mp.scorerhi)(candidate)
                    if not mp_fraction(written_interval.lower) < reference < mp_fraction(written_interval.upper):
                        raise ArithmeticError('mpmath disagrees with stored digits: ' + identity)
            numbers[identity] = dict(result, number=generated)
            records.append({'argument': identity,
                            'generator_interval_contained': True,
                            'quadrature_interval_contained': True,
                            'quadrature_accuracy_bits': independent.rel_accuracy_bits(),
                            'mpmath_decimal_precisions': [150, 200]})
            print(kind, identity, 'checked', flush=True)
        if stored is not None and set(stored) != set(numbers):
            raise ArithmeticError('stored and default-enumerated argument sets differ')
    evidence = {'table': generator.table, 'entries': len(numbers), 'digits': 100,
                'live': live, 'records': records,
                'versions': {package: importlib.metadata.version(package)
                             for package in ('python-flint', 'mpmath')}}
    evidence['source_sha256'] = {
        name: hashlib.sha256(Path(__file__).with_name(name).read_bytes()).hexdigest()
        for name in ('scorer.py', Path(__file__).name)}
    return numbers, evidence


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--kind', choices=('gi', 'hi'), required=True)
    parser.add_argument('--live', action='store_true')
    parser.add_argument('--output-dir', type=Path)
    arguments = parser.parse_args()
    numbers, evidence = check(arguments.kind, arguments.live)
    if arguments.output_dir:
        arguments.output_dir.mkdir(parents=True, exist_ok=True)
        (arguments.output_dir / 'numbers.json').write_text(json.dumps(numbers, indent=2) + '\n')
        (arguments.output_dir / 'checks-evidence.json').write_text(json.dumps(evidence, indent=2) + '\n')
    print(json.dumps({key: value for key, value in evidence.items() if key != 'records'}))