"""Recompute selected elementary, saddle-point and delay-equation values.

Install numberdb, python-flint and mpmath. Download this file beside the
table-specific generate.py and run `python generate.py` to verify, not write.

Arb encloses every elementary evaluation, exact rational argument and root
bracket. The delay solver keeps a polynomial and a separate uniform error
on each unit interval. If P has degree N and d_j are the coefficients of
P(y)/(a+y), then d_j=(p_j-d_(j-1))/a and, for j>N,
d_j=-d_(j-1)/a. Omitting from degree N therefore costs at most
|d_N| h^N/(1-h/a) on |y|<=h=1/2. An inherited uniform error E costs
E/(a-h) in the derivative. Integration from the left endpoint costs at
most the derivative error times the interval length, which is one.
The endpoint error is carried separately as well. Ball coefficients enclose
rounding errors; this proof does not replace a varying error by a constant
coefficient and then differentiate it.

The default 100-digit calculation uses 733 bits and degree 320. A second
run at higher precision/order, independent archived Sage implementations,
closed forms and quadrature controls are recorded in checks-evidence.json.
"""

import math
from contextlib import contextmanager
from fractions import Fraction
from functools import lru_cache

import numberdb
from flint import arb, ctx


SELECTED = {
    'T188': ('1',),
    'T191': ('2',) + tuple(str(10 ** exponent) for exponent in range(1, 23)),
    'T194': ('0', '1'),
    'T195': ('1',),
    'T205': ('0', '1'),
    'T206': ('0', 'dawson-critical-1', '1'),
    'T422': ('1', '2', '3', '5', '7', '9'),
    'T423': tuple(str(argument) for argument in range(1, 13)),
    'T433': tuple(str(argument) for argument in range(1, 11)),
    'T434': tuple(str(argument) for argument in range(1, 9)),
}
PARAMETER = {tid: ('s' if tid == 'T422' else 'u' if tid in
                  ('T423', 'T433', 'T434') else 'x') for tid in SELECTED}
GUARD_BITS = 400


@contextmanager
def precision(bits):
    previous = ctx.prec
    ctx.prec = bits
    try:
        yield
    finally:
        ctx.prec = previous


def as_ball(value):
    return value if isinstance(value, arb) else arb(str(value))


def endpoint(value):
    mantissa, exponent = value.man_exp()
    return Fraction(int(mantissa)) * Fraction(2) ** int(exponent)


def as_interval(value):
    if not value.is_finite():
        raise ArithmeticError('nonfinite ball')
    return numberdb.RealInterval(endpoint(value.lower()), endpoint(value.upper()))


def widen(value, radius):
    return value + arb(0, as_ball(radius).abs_upper())


def dawson(argument):
    return arb.pi().sqrt() * (-(argument * argument)).exp() * argument.erfi() / 2


def bracket_root(function, left, right, digits):
    left, right = Fraction(left), Fraction(right)
    left_sign = 1 if function(as_ball(left)) > 0 else -1
    if not function(as_ball(left)) * left_sign > 0:
        raise ArithmeticError('undecidable initial sign')
    if not function(as_ball(right)) * left_sign < 0:
        raise ArithmeticError('root is not bracketed')
    tolerance = Fraction(1, 10 ** (digits + 25))
    while right - left > tolerance:
        middle = (left + right) / 2
        value = function(as_ball(middle)) * left_sign
        if value > 0:
            left = middle
        elif value < 0:
            right = middle
        else:
            raise ArithmeticError('root sign lost; increase working precision')
    if not function(as_ball(left)) * left_sign > 0:
        raise ArithmeticError('left root endpoint failed')
    if not function(as_ball(right)) * left_sign < 0:
        raise ArithmeticError('right root endpoint failed')
    result = widen(as_ball((left + right) / 2), as_ball((right - left) / 2))
    if not function(result).contains(0):
        raise ArithmeticError('root residual does not contain zero')
    return result, left, right


@lru_cache(maxsize=12)
def dawson_root(digits, bits):
    with precision(bits):
        return bracket_root(lambda argument: 1 - 2 * argument * dawson(argument),
                            Fraction(1, 2), Fraction(1), digits)


def xi(argument, digits):
    if argument == 1:
        return 0
    return bracket_root(lambda root: root.exp() - 1 - argument * root,
                        1, 8, digits)[0]


def polynomial(coefficients, argument):
    result = arb(0)
    for coefficient in reversed(coefficients):
        result = result * argument + coefficient
    return result


class Model:
    def __init__(self, coefficients, error=0):
        self.coefficients = coefficients
        self.error = as_ball(error)

    def evaluate(self, argument):
        return widen(polynomial(self.coefficients, as_ball(argument)), self.error)


def constant_model(value, degree):
    return Model([as_ball(value)] + [arb(0) for index in range(degree)])


def integrate_step(previous, boundary, midpoint, sign, degree):
    half = arb(1) / 2
    midpoint = as_ball(midpoint)
    quotient = []
    prior = arb(0)
    for coefficient in previous.coefficients:
        prior = (coefficient - prior) / midpoint
        quotient.append(prior)
    omitted = abs(quotient[degree]) * half ** degree / (1 - half / midpoint)
    inherited = previous.error / (midpoint - half)
    current = [arb(0)] + [sign * quotient[index - 1] / index
                          for index in range(1, degree + 1)]
    current[0] = polynomial(boundary.coefficients, half) - polynomial(current, -half)
    return Model(current, boundary.error + omitted + inherited)


@lru_cache(maxsize=12)
def delay_models(kind, degree, bits):
    with precision(bits):
        if kind == 'rho':
            models = [constant_model(1, degree)]
            for interval in range(1, 10):
                models.append(integrate_step(models[-1], models[-1],
                                              Fraction(2 * interval + 1, 2), -1, degree))
            return models
        if kind == 'buchstab':
            models = [None, constant_model(1, degree)]
            for interval in range(2, 8):
                models.append(integrate_step(models[-1], models[-1],
                                              Fraction(2 * interval - 1, 2), 1, degree))
            return models
        if kind == 'sieve':
            upper, lower = [], []
            for interval in range(10):
                midpoint = Fraction(2 * interval - 1, 2)
                upper.append(constant_model(2 * arb.const_euler().exp(), degree)
                             if interval <= 2 else integrate_step(
                                 lower[interval - 1], upper[interval - 1], midpoint, 1, degree))
                lower.append(constant_model(0, degree) if interval <= 1 else integrate_step(
                    upper[interval - 1], lower[interval - 1], midpoint, 1, degree))
            return upper, lower
        raise ValueError('unknown delay model')


def delay_value(tid, identity, degree):
    argument = int(identity)
    half = Fraction(1, 2)
    if tid == 'T433':
        if argument == 1:
            return 1
        return delay_models('rho', degree, ctx.prec)[argument - 1].evaluate(half)
    if tid == 'T434':
        if argument <= 2:
            return Fraction(1, argument)
        return delay_models('buchstab', degree, ctx.prec)[argument - 1].evaluate(half) / argument
    return delay_models('sieve', degree, ctx.prec)[0][argument - 1].evaluate(half) / argument


def compute(tid, params, digits=100, extra_bits=0, degree=None):
    bits = numberdb.bits(digits, losing=GUARD_BITS + extra_bits)
    degree = degree or max(320, 3 * digits + 20)
    identity = str(params[PARAMETER[tid]])
    if identity not in SELECTED[tid]:
        raise ValueError('argument is outside the documented selection')
    with precision(bits):
        argument = (dawson_root(digits, bits)[0] if identity == 'dawson-critical-1'
                    else as_ball(Fraction(identity)))
        if tid == 'T188':
            value = argument.ei() if params['normalisation'] == 'Ei' else -(-argument).ei()
        elif tid == 'T191':
            value = argument.log().ei()
        elif tid == 'T194':
            value = argument.shi() if identity != '0' else 0
        elif tid == 'T195':
            value = argument.chi()
        elif tid == 'T205':
            value = argument.erfi() if identity != '0' else 0
        elif tid == 'T206':
            value = dawson(argument) if identity != '0' else 0
        elif tid == 'T423':
            value = xi(argument, digits)
        else:
            value = delay_value(tid, identity, degree)
        if isinstance(value, arb):
            if not value.is_finite() or value.rel_accuracy_bits() < math.ceil((digits + 12) * math.log2(10)):
                raise ArithmeticError('insufficient finite relative accuracy')
        return value


class RemainingValues(numberdb.Generator):
    type = 'R'
    digits = 100
    rigour = 'proven'
    files = ('generate.py', 'remaining_values.py', 'checks.py', 'checks-evidence.json',
             'relevance-audit.json')

    def __init__(self, tid):
        if tid not in SELECTED:
            raise ValueError('unassigned generator')
        self.table = tid
        self.parameters = (('normalisation', 'x') if tid == 'T188' else (PARAMETER[tid],))

    def enumerate(self):
        for normalisation in (('Ei', 'E1') if self.table == 'T188' else (None,)):
            for identity in SELECTED[self.table]:
                params = {PARAMETER[self.table]: identity}
                if normalisation is not None:
                    params = {'normalisation': normalisation, **params}
                yield params

    def value(self, params, digits):
        with precision(numberdb.bits(digits, losing=GUARD_BITS)):
            value = compute(self.table, params, digits)
            result = {'number': as_interval(value) if isinstance(value, arb) else value}
            if params.get('x') == 'dawson-critical-1':
                result.update({'param-latex': r'$x_{\mathrm{max}}$',
                               'comment': r'The unique positive maximum; its argument is defined in CITE{argument-selection}.'})
            return result