"""Certified Scorer values at intrinsically specified real arguments.

Requires python-flint >= 0.6, mpmath, and numberdb. No database access occurs
on import. Taylor coefficients satisfy c[k+3] = c[k]/((k+3)*(k+2)).
The three residue classes have decreasing absolute term ratios; each omitted
geometric tail is added to the radius. Root arguments remain balls throughout.

For Gi critical points, adjacent rational cells certify the root ordering.
On a cell with midpoint m and half-width h, Gronwall for (Gi, Gi') gives
B = (|Gi(m)| + |Gi'(m)| + h/pi)*exp(16*h) on |x| <= 16. Consequently
Gi lies in Gi(m) +/- h*B and Gi' in Gi'(m) +/- h*(R*B+1/pi).
Gi'' = x*Gi-1/pi then proves monotonicity where needed. Subdivision resolves
cells that neither exclude Gi'=0 nor prove monotonicity. Candidate roots
from mpmath are accepted only after independent ball sign certification.
"""

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

import mpmath as mp
import numberdb
from flint import arb, ctx


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


def working_precision(digits):
    if digits < 1:
        raise ValueError('digits must be positive')
    return math.ceil(digits * math.log2(10)) + 640


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


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


def exact_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 enclosure')
    return numberdb.RealInterval(exact_endpoint(value.lower()),
                                 exact_endpoint(value.upper()))


def sign(value):
    if not value.is_finite():
        raise ArithmeticError('nonfinite sign test')
    return 1 if value > 0 else -1 if value < 0 else 0


@lru_cache(maxsize=16)
def initial_coefficients(kind, precision):
    if precision != ctx.prec:
        raise ArithmeticError('coefficient precision mismatch')
    if kind not in ('gi', 'hi'):
        raise ValueError('unknown Scorer function')
    multiplier = 1 if kind == 'gi' else 2
    forcing = -1 if kind == 'gi' else 1
    return (multiplier / (arb(3) ** arb('7/6') * arb('2/3').gamma()),
            multiplier / (arb(3) ** arb('5/6') * arb('1/3').gamma()),
            forcing / (2 * arb.pi()))


def evaluate(argument, kind='gi', derivative=0, digits=100):
    if derivative not in (0, 1, 2):
        raise ValueError('only derivative orders zero through two')
    argument = as_ball(argument)
    radius = argument.abs_upper()
    tolerance = arb(10) ** (-digits - 80)
    answer = arb(0)
    for residue, coefficient in enumerate(initial_coefficients(kind, ctx.prec)):
        index = residue
        if index < derivative:
            coefficient /= (index + 3) * (index + 2)
            index += 3
        factor = math.factorial(index) // math.factorial(index - derivative)
        term = coefficient * factor * argument ** (index - derivative)
        subtotal = arb(0)
        while index < 20000:
            subtotal += term
            denominator = ((index + 3) * (index + 2),
                           index * (index + 2), index * (index - 1))[derivative]
            ratio = radius ** 3 / denominator
            if ratio < 1:
                tail = abs(term) * ratio / (1 - ratio)
                if tail < tolerance:
                    answer += widen(subtotal, tail)
                    break
            term *= argument ** 3 / denominator
            index += 3
        else:
            raise ArithmeticError('Taylor tail did not converge')
    if not answer.is_finite():
        raise ArithmeticError('nonfinite Taylor result')
    return answer


@lru_cache(maxsize=4096)
def point(argument, precision):
    if precision != ctx.prec:
        raise ArithmeticError('point precision mismatch')
    return evaluate(argument, digits=40), evaluate(argument, derivative=1, digits=40)


def local_enclosures(left, right, digits=40):
    midpoint = (left + right) / 2
    half_width = as_ball((right - left) / 2)
    if max(abs(left), abs(right)) > 16:
        raise ValueError('root-count certificate covers only |x| <= 16')
    if digits == 40:
        function, velocity = point(midpoint, ctx.prec)
    else:
        function = evaluate(midpoint, digits=digits)
        velocity = evaluate(midpoint, derivative=1, digits=digits)
    bound = (abs(function) + abs(velocity) + half_width / arb.pi())
    bound *= (16 * half_width).exp()
    function_box = widen(function, half_width * bound)
    velocity_box = widen(velocity, half_width * (
        as_ball(max(abs(left), abs(right))) * bound + 1 / arb.pi()))
    acceleration = widen(as_ball(midpoint), half_width) * function_box - 1 / arb.pi()
    return velocity_box, acceleration


def classify(left, right, depth=0):
    velocity, acceleration = local_enclosures(left, right)
    if sign(velocity):
        return []
    if sign(acceleration):
        left_sign = sign(point(left, ctx.prec)[1])
        right_sign = sign(point(right, ctx.prec)[1])
        if left_sign and right_sign:
            return [(left, right)] if left_sign != right_sign else []
    if depth >= 24:
        raise ArithmeticError('root count could not resolve a cell')
    midpoint = (left + right) / 2
    return classify(left, midpoint, depth + 1) + classify(midpoint, right, depth + 1)


@lru_cache(maxsize=1)
def critical_cells():
    with precision(working_precision(50)):
        positive = []
        for index in range(64):
            positive.extend(classify(Fraction(index, 64), Fraction(index + 1, 64)))
        negative = []
        for index in range(1024):
            negative.extend(reversed(classify(-Fraction(index + 1, 64),
                                              -Fraction(index, 64))))
            if len(negative) == 12:
                break
        if len(positive) != 1 or len(negative) != 12:
            raise ArithmeticError('unexpected number of critical points')
        return positive + negative


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


def critical_point(index, digits):
    if not 0 <= index <= 12:
        raise ValueError('certified Gi critical-point indices are 0 through 12')
    left, right = critical_cells()[index]
    places = digits + 80
    with mp.workdps(digits + 120):
        candidate = mp.findroot(lambda argument: mp.diff(mp.scorergi, argument),
                                (mp_fraction(left), mp_fraction(right)))
        center = Fraction(int(mp.nint(candidate * 10 ** places)), 10 ** places)
    epsilon = Fraction(1, 10 ** places)
    lower, upper = center - epsilon, center + epsilon
    if not left < lower < upper < right:
        raise ArithmeticError('candidate escaped its counted cell')
    if sign(evaluate(lower, derivative=1, digits=digits + 20)) * sign(
            evaluate(upper, derivative=1, digits=digits + 20)) != -1:
        raise ArithmeticError('critical point is not bracketed')
    if not sign(local_enclosures(lower, upper, digits=digits + 20)[1]):
        raise ArithmeticError('critical point is not isolated')
    return widen(as_ball(center), as_ball(epsilon))


def argument_for(identity, digits=100):
    if identity == '0':
        return arb(0)
    if identity.startswith('ai-zero-'):
        index = int(identity[len('ai-zero-'):])
        if index < 1:
            raise ValueError('Airy zero index starts at one')
        return arb.airy_ai_zero(index)
    if identity.startswith('gi-critical-'):
        return critical_point(int(identity[len('gi-critical-'):]), digits)
    return as_ball(Fraction(identity))


class ScorerValues(numberdb.Generator):
    parameters = ('x',)
    type = 'R'
    digits = 100
    rigour = 'proven'
    files = ('generate.py', 'scorer.py', 'checks.py', 'selection.md', 'checks-evidence.json')

    def enumerate(self, airy_limit=20, critical_limit=12):
        if airy_limit < 0 or not 0 <= critical_limit <= 12:
            raise ValueError('invalid enumeration bounds')
        yield {'x': '0'}
        for index in range(1, airy_limit + 1):
            yield {'x': 'ai-zero-%d' % index}
        if self.kind == 'gi':
            for index in range(critical_limit + 1):
                yield {'x': 'gi-critical-%d' % index}

    def value(self, params, digits):
        identity = str(params['x'])
        if self.kind == 'hi' and identity.startswith('gi-critical-'):
            raise ValueError('Gi critical points are not default Hi arguments')
        for attempt in range(4):
            with precision(working_precision(digits) * 2 ** attempt):
                argument = argument_for(identity, digits)
                value = evaluate(argument, self.kind, digits=digits)
                if value.rel_accuracy_bits() < math.ceil((digits + 10) * math.log2(10)):
                    continue
                result = {'number': as_interval(value)}
                if identity.startswith('ai-zero-'):
                    index = int(identity[len('ai-zero-'):])
                    result['param-latex'] = r'$a_{%d}$' % index
                    result['comment'] = (
                        r'$a_{%d}$ is the negative zero of $\operatorname{Ai}$ with index $n=%d$, '
                        r'counted from the origin. See CITE{brownian-application}.'
                    ) % (index, index)
                elif identity.startswith('gi-critical-'):
                    index = int(identity[len('gi-critical-'):])
                    result['param-latex'] = r'$\xi_{%d}$' % index
                    result['comment'] = (
                        r'Extremal height, with $\operatorname{Gi}^{\prime}(\xi_{%d})=0$; '
                        r'the root convention is CITE{critical-arguments}.'
                    ) % index
                elif identity == '0':
                    result['comment'] = 'The distinguished origin; see CITE{initial-values}.'
                return result
        raise ArithmeticError('insufficient certified digits after precision escalation')