"""Independent, read-only checks of the selected Lambert W values.

    pip install numberdb python-flint mpmath
    python t200_checks.py --document table.json

Without --document, fetch the current table using numberdb. Every nonexact
value is checked by a sign-certified root bracket that does not call Lambert
W, and by mpmath at 150 and 200 digits. Both complete Arb enclosures must lie
inside the written decimal's last-place interval. Irrational arguments are
recomputed from exact expressions in both libraries.

For the Planck peaks z=m(1-exp(-z)), the bracket (m-1,m) excludes z=0.
For the 3-uniform 2-core, exp(lambda)-1=2 lambda has a unique positive root
in (1,2); its transformation is W_-1(-exp(-1/2)/2)=-lambda-1/2.
The exponential derivative is positive throughout (1,2). The Planck
equation has negative second derivative, one positive maximum, and only one
positive zero for m>1. Omega is the unique zero of w exp(w)-1 in (0,1).
"""

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 t200_values import (EXACT, GUARD_BITS, SELECTION, LambertValues, argument,
                         as_interval, compute, precision)


def written_bounds(text):
    text = str(text)
    if '.' not in text and 'e' not in text.lower():
        exact = Fraction(text)
        return exact, exact
    decimal = Decimal(text)
    center = Fraction(decimal)
    unit = Fraction(10) ** decimal.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 in ('0', '1'):
        return mp.mpf(identity)
    if identity == 'e':
        return mp.exp(1)
    if identity == '-1/e':
        return -mp.exp(-1)
    if identity == '-3*exp(-3)':
        return -3 * mp.exp(-3)
    if identity == '-5*exp(-5)':
        return -5 * mp.exp(-5)
    if identity == '-exp(-1/2)/2':
        return -mp.exp(-mp.mpf(1) / 2) / 2
    raise ValueError(identity)


def sign(value):
    if not value.is_finite():
        raise ArithmeticError('nonfinite bracket evaluation')
    if value > 0:
        return 1
    if value < 0:
        return -1
    raise ArithmeticError('bracket endpoint sign is undecidable')


def root_bracket(function, left, right, digits):
    left, right = Fraction(left), Fraction(right)
    left_sign = sign(function(arb(str(left))))
    if sign(function(arb(str(right)))) != -left_sign:
        raise ArithmeticError('initial interval does not bracket a root')
    tolerance = Fraction(1, 10 ** (digits + 25))
    while right - left > tolerance:
        middle = (left + right) / 2
        if sign(function(arb(str(middle)))) == left_sign:
            left = middle
        else:
            right = middle
    if (sign(function(arb(str(left)))) != left_sign or
            sign(function(arb(str(right)))) != -left_sign):
        raise ArithmeticError('final root bracket failed')
    return left, right


def independent_bracket(identity, digits=100):
    if identity == '1':
        function = lambda point: point * point.exp() - 1
        lower, upper = root_bracket(function, 0, 1, digits)
        result = lower, upper
        equation = 'w*exp(w)=1; root in (0,1)'
    elif identity in ('-3*exp(-3)', '-5*exp(-5)'):
        order = 3 if identity == '-3*exp(-3)' else 5
        function = lambda point: order * (1 - (-point).exp()) - point
        lower, upper = root_bracket(function, order - 1, order, digits)
        result = lower - order, upper - order
        equation = 'z=%d*(1-exp(-z)); positive root in (%d,%d)' % (order, order - 1, order)
    elif identity == '-exp(-1/2)/2':
        function = lambda point: point.exp() - 1 - 2 * point
        lower, upper = root_bracket(function, 1, 2, digits)
        result = -upper - Fraction(1, 2), -lower - Fraction(1, 2)
        equation = 'exp(lambda)-1=2*lambda; positive root in (1,2)'
    else:
        raise ValueError(identity)
    return result, {'equation': equation, 'root_lower': str(lower),
                    'root_upper': str(upper), 'opposite_endpoint_signs_certified': True,
                    'value_lower': str(result[0]), 'value_upper': str(result[1])}


def verify_document(document):
    with precision(numberdb.bits(150, losing=GUARD_BITS)):
        control = root_bracket(lambda point: point - 1, 0, 3, 100)
        if not control[0] < 1 < control[1]:
            raise ArithmeticError('known-root control failed')
    generator = LambertValues()
    numbers = document['Numbers']
    if list(document['Parameters']) != list(generator.parameters):
        raise AssertionError('parameter order changed')
    if list(numbers) != list(SELECTION):
        raise AssertionError('branch enumeration differs')
    for branch, identities in SELECTION.items():
        if list(numbers[branch]) != list(identities):
            raise AssertionError('stored arguments differ from generator enumeration')
    records = []
    for params in generator.enumerate():
        branch, identity = params['branch'], params['x']
        entry = numbers[branch][identity]
        written = entry['number'] if isinstance(entry, dict) else entry
        lower, upper = written_bounds(written)
        record = {'params': params, 'written': written}
        if (branch, identity) in EXACT:
            expected = EXACT[(branch, identity)]
            if lower != expected or upper != expected or written != str(expected):
                raise AssertionError('exact inverse identity was not stored exactly')
            record['exact_proof'] = 'W_k(y*exp(y))=y on the specified real branch'
            with precision(numberdb.bits(100, losing=GUARD_BITS)):
                residual = arb(expected) * arb(expected).exp() - argument(identity)
                if not residual.is_finite() or not residual.contains(0):
                    raise ArithmeticError('exact identity consistency control failed')
            records.append(record)
            continue
        significant_digits = len(Decimal(written).as_tuple().digits)
        if significant_digits < 100:
            raise ArithmeticError('stored approximation has fewer than 100 significant digits')
        record['significant_digits'] = significant_digits
        records_arb = []
        for digits, extra_bits in ((100, 0), (130, 80)):
            bits = numberdb.bits(digits, losing=GUARD_BITS + extra_bits)
            with precision(bits):
                value = compute(branch, identity, digits, extra_bits)
                interval = as_interval(value)
                if not lower <= interval.lower <= interval.upper <= upper:
                    raise ArithmeticError('Arb enclosure is outside the written interval')
                argument_ball = argument(identity)
                records_arb.append({'working_bits': bits,
                                    'value_relative_accuracy_bits': value.rel_accuracy_bits(),
                                    'argument_relative_accuracy_bits': argument_ball.rel_accuracy_bits(),
                                    'value_lower': str(interval.lower),
                                    'value_upper': str(interval.upper),
                                    'entire_enclosure_contained': True})
        with precision(numberdb.bits(150, losing=GUARD_BITS)):
            bracket, evidence = independent_bracket(identity)
            if not lower <= bracket[0] < bracket[1] <= upper:
                raise ArithmeticError('independent root bracket is outside written interval')
            tight = as_interval(compute(branch, identity, 130, 80))
            if not max(bracket[0], tight.lower) <= min(bracket[1], tight.upper):
                raise ArithmeticError('certified independent methods disagree')
            if branch == '0' and not bracket[0] > -1:
                raise ArithmeticError('root bracket has wrong branch')
            if branch == '-1' and not bracket[1] < -1:
                raise ArithmeticError('root bracket has wrong branch')
        independent = []
        for digits in (150, 200):
            with mp.workdps(digits):
                value = mp.lambertw(mp_argument(identity), int(branch))
                if (not mp.isfinite(value) or not mp.im(value) == 0 or
                        not mp_fraction(lower) <= value <= mp_fraction(upper)):
                    raise ArithmeticError('mpmath disagreement')
                independent.append({'digits': digits, 'value': mp.nstr(value, 110)})
        record.update({'arb_enclosures': records_arb, 'root_bracket': evidence,
                       'mpmath': independent})
        records.append(record)
    return {'table': 'T200', 'entries': len(records), 'digits': 100,
            'known_root_control': {'equation': 'w-1=0', 'initial_bracket': [0, 3],
                                   'exact_answer': 1, 'passed': True},
            'exact_entries': len(EXACT), 'nonexact_entries': len(records) - len(EXACT),
            'every_written_interval_verified': True,
            'argument_uncertainty_propagated': True, 'records': records,
            'versions': {name: importlib.metadata.version(name)
                         for name in ('python-flint', 'mpmath')}}


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--document', type=Path)
    parser.add_argument('--output', type=Path)
    options = parser.parse_args()
    document = (json.loads(options.document.read_text()) if options.document
                else numberdb.table('T200'))
    result = verify_document(document)
    encoded = json.dumps(result, indent=2) + '\n'
    if options.output:
        options.output.write_text(encoded)
    print('T200: %d exact identities and %d certified, independently checked values' %
          (result['exact_entries'], result['nonexact_entries']))
