"""Independently check a prepared table without API writes.

    pip install numberdb python-flint mpmath
    python balanced_checks.py T200 --document table.json --selection selection.json

Arb enclosures must fit within the exact rational endpoints denoted by every
stored decimal. Independent mpmath quadrature, Lambert inversion and closed
forms check every new value at 150 and 200 decimal digits. Previously stored
deep delay values retain their provenance; their supplementary check here is
higher-order enclosure agreement, not an independent full recalculation.
"""

import argparse
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 arb

from balanced_values import (GUARD_BITS, PARAMETERS, argument, as_ball,
                             as_interval, compute, dawson, dawson_root, precision)


def written_bounds(text):
    text = str(text)
    if '/' in text or ('.' not in text and 'e' not in text.lower()):
        exact = Fraction(text)
        return exact, exact
    center = Fraction(Decimal(text))
    unit = Fraction(10) ** Decimal(text).as_tuple().exponent
    return center - unit, center + unit


def mp_fraction(value):
    return mp.mpf(value.numerator) / value.denominator


def mp_argument(identity):
    if identity == 'e':
        return mp.e
    if identity == '-1/e':
        return -1 / mp.e
    if identity.startswith(('log(', 'exp(', 'sqrt(')):
        operation, integer = identity[:-1].split('(')
        return getattr(mp, operation)(int(integer))
    if '*exp(' in identity:
        coefficient = mp_fraction(Fraction(identity.split('*')[0]))
        return coefficient * mp.exp(coefficient)
    if identity.startswith('-exp('):
        denominator = int(identity.rsplit('/', 1)[1])
        return -mp.exp(-mp.mpf(1) / denominator) / denominator
    if identity == 'dawson-critical-1':
        return mp.findroot(lambda point: 1 - 2 * point * mp.exp(-point ** 2) *
                           mp.quad(lambda variable: mp.exp(variable ** 2), [0, point]),
                           (mp.mpf('0.9'), 1))
    return mp_fraction(Fraction(identity))


def rho_initial(point):
    if point <= 1:
        return mp.mpf(1)
    if point <= 2:
        return 1 - mp.log(point)
    if point <= 3:
        return (1 - mp.log(point) + mp.log(point - 1) * mp.log(point)
                + mp.polylog(2, 1 - point) + mp.pi ** 2 / 12)
    return rho_initial(mp.mpf(3)) - mp.quad(
        lambda variable: rho_initial(variable - 1) / variable, [3, point])


def independent_value(tid, params):
    point = mp_argument(params[PARAMETERS[tid][-1]])
    if tid == 'T188':
        if params['normalisation'] == 'E1':
            return mp.exp(-point) * mp.quad(lambda variable: mp.exp(-variable) /
                                           (point + variable), [0, 1, mp.inf])
        return mp.euler + mp.log(point) + mp.quad(
            lambda variable: mp.expm1(variable) / variable if variable else 1, [0, point])
    if tid == 'T191':
        return mp.li(point)
    if tid == 'T194':
        return mp.quad(lambda variable: mp.sinh(variable) / variable if variable else 1,
                       [0, point])
    if tid == 'T195':
        return mp.euler + mp.log(point) + mp.quad(
            lambda variable: 2 * mp.sinh(variable / 2) ** 2 / variable if variable else 0,
            [0, point])
    if tid in ('T205', 'T206'):
        integral = mp.quad(lambda variable: mp.exp(variable ** 2), [0, point])
        return 2 * integral / mp.sqrt(mp.pi) if tid == 'T205' else mp.exp(-point ** 2) * integral
    if tid == 'T200':
        exact = {('0', '0'): 0, ('0', 'e'): 1,
                 ('0', '-1/e'): -1, ('-1', '-1/e'): -1}
        key = (params['branch'], params['x'])
        return mp.mpf(exact[key]) if key in exact else mp.lambertw(point, int(params['branch']))
    if tid == 'T423':
        return mp.mpf(0) if point == 1 else -mp.lambertw(-mp.exp(-1 / point) / point, -1) - 1 / point
    if tid == 'T433' and point <= 4:
        return rho_initial(point)
    if tid == 'T434' and point <= 4:
        if point <= 2:
            return 1 / point
        if point <= 3:
            return (1 + mp.log(point - 1)) / point
        return (1 + mp.log(2) + mp.quad(
            lambda variable: (1 + mp.log(variable - 2)) / (variable - 1), [3, point])) / point
    if tid == 'T422' and point <= 5:
        return 2 * mp.exp(mp.euler) / point * (1 + (mp.quad(
            lambda variable: mp.log(variable - 2) / (variable - 1), [3, point])
            if point > 3 else 0))
    return None


def leaf(numbers, params):
    value = numbers
    for identity in params.values():
        value = value[str(identity)]
    return value


def rows_of(document):
    parameters = tuple(document['Parameters'])

    def descend(node, depth, params):
        if depth == len(parameters):
            yield params, node
        else:
            for identity, child in node.items():
                yield from descend(child, depth + 1, {**params, parameters[depth]: str(identity)})

    return list(descend(document['Numbers'], 0, {}))


def verify_document(tid, document, selection):
    with mp.workdps(150):
        control = mp.quad(lambda variable: mp.exp(variable), [0, 1])
        if abs(control - mp.expm1(1)) > mp.mpf('1e-140') or not control > 1:
            raise ArithmeticError('nonzero quadrature control failed')
        if abs(mp.lambertw(mp.e) - 1) > mp.mpf('1e-140'):
            raise ArithmeticError('Lambert inverse control failed')
        if abs(rho_initial(mp.mpf(2)) - (1 - mp.log(2))) > mp.mpf('1e-140'):
            raise ArithmeticError('delay initial-segment control failed')
    if [params for params, entry in rows_of(document)] != [row['params'] for row in selection['rows']]:
        raise ArithmeticError('enumeration differs from complete document selection')
    records = []
    for record in selection['rows']:
        params = record['params']
        entry = leaf(document['Numbers'], params)
        written = entry['number'] if isinstance(entry, dict) else entry
        lower, upper = written_bounds(written)
        for digits, extra_bits, degree in ((100, 0, 320), (130, 80, 420)):
            value = compute(tid, params, digits, extra_bits, degree)
            with precision(numberdb.bits(digits, losing=GUARD_BITS + extra_bits)):
                interval = as_interval(value) if isinstance(value, arb) else numberdb.RealInterval(value, value)
                if not lower <= interval.lower <= interval.upper <= upper:
                    raise ArithmeticError(tid + ': enclosure outside written interval: ' + str(params))
                accuracy = value.rel_accuracy_bits() if isinstance(value, arb) else None
        comparisons = []
        for digits in (150, 200):
            with mp.workdps(digits):
                independent = independent_value(tid, params)
                if independent is not None:
                    if not mp.isfinite(independent) or not mp_fraction(lower) <= independent <= mp_fraction(upper):
                        raise ArithmeticError(tid + ': independent mismatch: ' + str(params))
                    comparisons.append(digits)
        if not record['retained'] and not comparisons:
            raise ArithmeticError('new value has no independent check')
        bracket_check = False
        with precision(900):
            if tid == 'T200' and lower != upper:
                source = argument(params['x'])
                sign = 1 if params['branch'] == '0' else -1
                if not ((as_ball(lower) * as_ball(lower).exp() - source) * sign < 0 and
                        (as_ball(upper) * as_ball(upper).exp() - source) * sign > 0):
                    raise ArithmeticError('Lambert written interval does not bracket root')
                bracket_check = True
            if tid == 'T423' and lower != upper:
                source = argument(params['u'])
                if not (as_ball(lower).exp() - 1 - source * as_ball(lower) < 0 and
                        as_ball(upper).exp() - 1 - source * as_ball(upper) > 0):
                    raise ArithmeticError('saddle-root written interval does not bracket root')
                bracket_check = True
        records.append({'params': params, 'retained': record['retained'],
                        'written': written, 'enclosure_inside_written_interval': True,
                        'higher_precision_order_check': True,
                        'higher_run_relative_accuracy_bits': accuracy,
                        'independent_mpmath_precisions': comparisons,
                        'written_interval_root_bracket': bracket_check})
    evidence = {'table': tid, 'records': records, 'all_new_values_independently_checked': True,
                'counts': {'total': len(records), 'new': sum(not row['retained'] for row in records)},
                'controls': ['integral exp(t), 0..1 = e-1, nonzero', 'W_0(e)=1', 'rho(2)=1-log(2)'],
                'working_bits': numberdb.bits(100, losing=GUARD_BITS),
                'second_working_bits': numberdb.bits(130, losing=GUARD_BITS + 80),
                'delay_degrees': [320, 420],
                'versions': {name: importlib.metadata.version(name) for name in ('python-flint', 'mpmath')},
                'limitations': ['Retained delay values beyond the initial branches have enclosure and higher-order checks here; their earlier independent provenance remains attached.']
                if tid in ('T422', 'T433', 'T434') else []}
    if tid == 'T206':
        with precision(900):
            root, left, right = dawson_root(130, 900)
            if not (1 - 2 * as_ball(left) * dawson(as_ball(left)) > 0 and
                    1 - 2 * as_ball(right) * dawson(as_ball(right)) < 0):
                raise ArithmeticError('Dawson derivative bracket failed')
            if not (dawson(root) - 1 / (2 * root)).contains(0):
                raise ArithmeticError('Dawson maximal-height identity failed')
            evidence['critical_argument'] = {'left_exact': str(left), 'right_exact': str(right),
                                            'strict_derivative_signs': [1, -1],
                                            'argument_error_propagated': True}
    return evidence


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('tid', choices=PARAMETERS)
    parser.add_argument('--document', type=Path, required=True)
    parser.add_argument('--selection', type=Path, required=True)
    arguments = parser.parse_args()
    print(json.dumps(verify_document(arguments.tid,
                                    json.loads(arguments.document.read_text()),
                                    json.loads(arguments.selection.read_text())), indent=2))
