"""Regular Coulomb wave function values -- numberdb.org/T510.

Run with SageMath (install the NumberDB client with sage -pip install numberdb):
    sage -python generate.py             # verify all stored entries
    sage -python generate.py --fill      # fill the existing draft using a key
    sage -python checks.py               # independent local checks
    sage -python checks.py --stored      # checks on values read from the draft

The account must be allowed to read the draft. Set NUMBERDB_API_KEY in the
calling environment, or pipe the key on stdin with NUMBERDB_KEY_FROM_STDIN=1.
Filling entries does not publish or offer the table for review.

DLMF 33.2.4 and 33.2.5 fix the normalization. All arithmetic is Sage/Arb
complex-ball arithmetic, including the gamma normalization and rational
arguments. Fifty extra decimal working digits are used. Every returned ball
must have at least ten more supported decimal digits than the requested 100.
The measured accuracy and independent checks are recorded by checks.py.

The 88 arguments follow family #210: ell=0..3; eta=0,+/-1/2,+/-1,+/-2;
rho=1/2,1,2, plus rho=5 at eta=2. Positive and negative eta are adjacent.
"""
import os
import sys
from math import factorial

import numberdb.sage as numberdb
from sage.rings.complex_arb import ComplexBallField
from sage.rings.rational_field import QQ

EXTRA_DECIMAL_DIGITS = 50
ETA = ('0', '1/2', '-1/2', '1', '-1', '2', '-2')
RADII = ('1/2', '1', '2')


def regular_ball(ell, eta, rho, digits=100):
    """Enclose DLMF F_ell(eta,rho), never a rounded fixed-precision answer."""
    ell = int(ell)
    eta, rho = QQ(eta), QQ(rho)
    if ell < 0 or rho <= 0:
        raise ValueError('ell must be nonnegative and rho positive')
    C = ComplexBallField(numberdb.bits(digits + EXTRA_DECIMAL_DIGITS))
    e, r, i = C(eta), C(rho), C(0, 1)
    a, b = ell + 1 - i * e, C(2 * ell + 2)
    normal = (C(2)**ell * (-C.pi() * e / 2).exp()
              * abs((ell + 1 + i * e).gamma()) / factorial(2 * ell + 1))
    result = normal * r**(ell + 1) * (-i * r).exp() \
        * (2 * i * r).hypergeometric([a], [b])
    if (not result.real().is_finite() or not result.imag().is_finite()) or not result.imag().contains_zero():
        raise ArithmeticError('Coulomb F did not return a finite real enclosure')
    value = result.real()
    if value.accuracy() < numberdb.bits(digits + 10):
        raise ArithmeticError('The Coulomb enclosure has insufficient precision')
    return value


class RegularCoulomb(numberdb.Generator):
    table = 'T510'
    parameters = ('ell', 'eta', 'rho')
    type = 'R'
    digits = 100
    rigour = 'proven'
    files = ('generate.py', 'checks.py', 'check.py')

    def enumerate(self):
        for ell in range(4):
            for eta in ETA:
                radii = RADII + (('5',) if eta == '2' else ())
                for rho in radii:
                    yield {'ell': str(ell), 'eta': eta, 'rho': rho}

    def value(self, params, digits):
        return regular_ball(params['ell'], params['eta'], params['rho'], digits)


if __name__ == '__main__':
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN') == '1':
        os.environ['NUMBERDB_API_KEY'] = sys.stdin.read().strip()
    generator = RegularCoulomb()
    if '--fill' in sys.argv or os.environ.get('NUMBERDB_PUBLISH') == '1':
        print(generator.publish(message='Fill 88 rigorously enclosed Coulomb F values; attach generator and checks'))
    else:
        report = generator.verify(sample=None)
        print(report)
        sys.exit(0 if report.ok else 1)
