"""Certified Kelvin/Struve extrema and Coulomb turning-point values.

Install: python -m pip install numberdb python-flint mpmath
Use the attached generate.py to verify its table without writing.

Critical arguments are defined by an isolating rational interval, not a
rounded decimal or a claimed global root ordinal. Opposite endpoint signs
give existence. Gronwall bounds and the differential equation prove that the
second derivative has a fixed nonzero sign throughout the isolating interval,
giving uniqueness and a genuine extremum. No completeness claim is made
outside these intervals. mpmath proposes a narrow bracket but never certifies
it. Arb certifies both its endpoints and the propagated argument uncertainty.
"""

import json
import math
from contextlib import contextmanager
from fractions import Fraction
from functools import lru_cache
from pathlib import Path

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


KINDS = {'T490': 'ber', 'T491': 'bei', 'T492': 'struveh', 'T493': 'ker',
         'T494': 'struvel', 'T509': 'kei', 'T510': 'coulombf', 'T511': 'coulombg'}
GUARD_BITS = 440


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


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


def interval(lower, upper):
    lower, upper = Fraction(lower), Fraction(upper)
    return arb(str((lower + upper) / 2), str((upper - lower) / 2))


def struve(kind, order, argument):
    order = arb(str(order))
    half = arb('1/2')
    scaled = acb(argument * argument / 4)
    if kind == 'struveh':
        scaled = -scaled
    result = ((acb(argument) / 2) ** (order + 1)
              * scaled.hypgeom([1], [3 * half, order + 3 * half])
              / ((3 * half).gamma() * (order + 3 * half).gamma()))
    if not result.is_finite() or not result.imag.contains(0):
        raise ArithmeticError('Struve result is not a finite real enclosure')
    return result.real


def kelvin_pair(kind, argument):
    rotation = acb(-1 if kind in ('ber', 'bei') else 1, 1) / arb(2).sqrt()
    rotated = argument * rotation
    if kind in ('ber', 'bei'):
        value = rotated.bessel_j(0)
        derivative = -rotation * rotated.bessel_j(1)
    else:
        value = rotated.bessel_k(0)
        derivative = -rotation * rotated.bessel_k(1)
    return value, derivative


def values_and_derivatives(kind, order, argument):
    if kind in ('ber', 'bei', 'ker', 'kei'):
        if int(order) != 0:
            raise ValueError('This selection contains only order-zero Kelvin functions')
        value, derivative = kelvin_pair(kind, argument)
        second = acb(0, 1) * value - derivative / argument
        component = 'real' if kind in ('ber', 'ker') else 'imag'
        return tuple(getattr(item, component) for item in (value, derivative, second))
    if kind != 'struveh':
        raise ValueError('No stationary-point family implemented for ' + kind)
    value = struve(kind, order, argument)
    derivative = struve(kind, order - 1, argument) - order * value / argument
    forcing = (argument / 2) ** (order - 1) / (
        arb.pi().sqrt() * arb(str(Fraction(order) + Fraction(1, 2))).gamma())
    second = forcing - derivative / argument - (1 - order * order / argument**2) * value
    return value, derivative, second


def second_derivative_on_interval(kind, order, lower, upper):
    domain = interval(lower, upper)
    middle = arb(str((Fraction(lower) + Fraction(upper)) / 2))
    half_width = arb(str((Fraction(upper) - Fraction(lower)) / 2))
    left = arb(str(lower))
    if kind in ('ber', 'bei', 'ker', 'kei'):
        value, derivative = kelvin_pair(kind, middle)
        components = (value.real, value.imag, derivative.real, derivative.imag)
        norm = max(component.abs_upper() for component in components)
        matrix_bound = 1 + 1 / left
        variation = half_width * matrix_bound * norm * (matrix_bound * half_width).exp()
        enclosed = [component + arb(0, variation.abs_upper()) for component in components]
        if kind in ('ber', 'ker'):
            return -enclosed[2] / domain - enclosed[1]
        return -enclosed[3] / domain + enclosed[0]
    value, derivative, unused = values_and_derivatives(kind, order, middle)
    matrix_bound = 1 + 1 / left + order * order / left**2
    forcing = (domain / 2) ** (order - 1) / (
        arb.pi().sqrt() * arb(str(Fraction(order) + Fraction(1, 2))).gamma())
    norm = max(value.abs_upper(), derivative.abs_upper())
    variation = half_width * (matrix_bound * norm + forcing.abs_upper()) * (
        matrix_bound * half_width).exp()
    value += arb(0, variation.abs_upper())
    derivative += arb(0, variation.abs_upper())
    return forcing - derivative / domain - (1 - order * order / domain**2) * value


def mp_value(kind, order, argument):
    return getattr(mp, kind)(order, argument)


def mp_derivative(kind, order, argument):
    if kind in ('ber', 'bei', 'ker', 'kei'):
        rotation = mp.exp((3 if kind in ('ber', 'bei') else 1) * mp.pi * 1j / 4)
        bessel = mp.besselj if kind in ('ber', 'bei') else mp.besselk
        derivative = -rotation * bessel(1, argument * rotation)
        return derivative.real if kind in ('ber', 'ker') else derivative.imag
    return mp.struveh(order - 1, argument) - order * mp.struveh(order, argument) / argument


@lru_cache(maxsize=256)
def critical_argument(kind, order, lower, upper, bits):
    if bits != ctx.prec:
        raise ArithmeticError('Root cache precision mismatch')
    endpoint_signs = [sign(values_and_derivatives(kind, order, arb(endpoint))[1])
                      for endpoint in (lower, upper)]
    if endpoint_signs[0] * endpoint_signs[1] != -1:
        raise ArithmeticError('Root interval lacks opposite endpoint signs')
    second = second_derivative_on_interval(kind, order, lower, upper)
    if sign(second) == 0:
        raise ArithmeticError('Root uniqueness not proved on isolating interval')
    decimal_precision = math.floor(bits * math.log10(2))
    with mp.workdps(70):
        endpoints = [mp.mpf(str(Fraction(endpoint).numerator)) /
                     Fraction(endpoint).denominator for endpoint in (lower, upper)]
        candidate = mp.findroot(lambda argument: mp_derivative(kind, order, argument), endpoints,
            solver='anderson')
        approximation = Fraction(mp.nstr(candidate, 65))
    for iteration in range(5):
        point = arb(str(approximation))
        unused, derivative, curvature = values_and_derivatives(kind, order, point)
        approximation = exact_endpoint((point - derivative / curvature).mid())
    radius = Fraction(1, 10**(decimal_precision - 60))
    narrow_lower = approximation - radius
    narrow_upper = approximation + radius
    if not Fraction(lower) < narrow_lower < narrow_upper < Fraction(upper):
        raise ArithmeticError('Candidate escaped its defining interval')
    narrow_signs = [sign(values_and_derivatives(kind, order, arb(str(endpoint)))[1])
                    for endpoint in (narrow_lower, narrow_upper)]
    if narrow_signs != endpoint_signs:
        raise ArithmeticError('Uncertified narrow bracket')
    argument = interval(narrow_lower, narrow_upper)
    return argument, {'lower': str(narrow_lower), 'upper': str(narrow_upper),
                      'endpoint_signs': narrow_signs,
                      'second_derivative_sign': sign(second),
                      'defining_interval': [lower, upper]}


def evaluate(kind, params, argument):
    if kind in ('ber', 'bei', 'ker', 'kei'):
        if argument == 0:
            return {'ber': 1, 'bei': 0, 'kei': -arb.pi() / 4}[kind]
        value, unused = kelvin_pair(kind, argument)
        return value.real if kind in ('ber', 'ker') else value.imag
    if kind in ('struveh', 'struvel'):
        if argument == 0:
            return 0
        return struve(kind, int(params['nu']), argument)
    result = getattr(acb(argument), 'coulomb_f' if kind == 'coulombf' else 'coulomb_g')(
        int(params['ell']), acb(arb(params['eta'])))
    if not result.is_finite() or not result.imag.contains(0):
        raise ArithmeticError('Coulomb result is not a finite real enclosure')
    return result.real


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

    def __init__(self, tid, selection_path=None):
        self.table = tid
        self.kind = KINDS[tid]
        self.parameters = ('ell', 'eta', 'rho') if tid in ('T510', 'T511') else ('nu', 'x')
        path = Path(selection_path) if selection_path else Path(__file__).with_name('selection.json')
        self.selection = json.loads(path.read_text())
        if self.selection['tid'] != tid:
            raise ValueError('Selection belongs to another table')

    def enumerate(self):
        for row in self.selection['rows']:
            yield row['params']

    def argument(self, params):
        row = next(row for row in self.selection['rows'] if row['params'] == params)
        if 'bracket' in row:
            return critical_argument(self.kind, int(params['nu']), *row['bracket'], ctx.prec)
        if row['argument'] == 'turning':
            coupling = arb(params['eta'])
            order = int(params['ell'])
            return coupling + (coupling**2 + order * (order + 1)).sqrt(), None
        return arb(row['argument']), None

    def value(self, params, digits):
        with precision(math.ceil(digits * math.log2(10)) + GUARD_BITS):
            argument, unused = self.argument(params)
            result = evaluate(self.kind, params, argument)
            if isinstance(result, int):
                return result
            if not result.is_finite() or result.rel_accuracy_bits() < math.ceil((digits + 15) * math.log2(10)):
                raise ArithmeticError('Insufficient certified accuracy')
            return as_interval(result)


def verify(tid):
    generator = IntrinsicValues(tid)
    report = generator.verify(sample=None)
    print(report)
    raise SystemExit(0 if report.ok else 1)
