"""Irregular Coulomb wave function values -- numberdb.org/T511.

Run with SageMath (or passagemath with the required Arb modules):
    sage -pip install numberdb
    sage -python generate.py                 # verify all stored entries
    sage -python checks.py                  # numerical checks, no writes
    sage -python generate.py --fill-draft    # fill the existing draft

The fill command requires NUMBERDB_API_KEY; it does not publish the table or
request review. The campaign also permits a key on stdin when
NUMBERDB_KEY_FROM_STDIN=1 and filling when NUMBERDB_PUBLISH=1.

DLMF 33.2.7 gives H+ = (-i)^ell exp(pi*eta/2) Gamma(ell+1+i*eta)/
abs(Gamma(ell+1+i*eta)) exp(i*rho) (-2*i*rho)^(ell+1)
U(ell+1+i*eta,2*ell+2,-2*i*rho). H- is its conjugate for real arguments,
and G=(H+ + H-)/2. Every operation, including rational conversion, is
performed in Sage/Arb complex balls. Principal branches are used.

The initial working precision is numberdb.bits(digits+50), 515 bits at
100 requested digits. The returned real ball must have at least
numberdb.bits(digits+10) bits of relative accuracy; otherwise precision is
doubled, with a finite limit. The finite selection is specified by the
Coulomb family proposal: four partial waves, seven simple reference couplings,
three radii, and the handbook point rho=5 only at eta=2.
"""
import numberdb.sage as numberdb
from sage.rings.complex_arb import ComplexBallField
from sage.rings.rational_field import QQ
import os
import sys

WORKING_EXTRA_DIGITS = 50
COUPLINGS = ('0', '1/2', '-1/2', '1', '-1', '2', '-2')
RADII = ('1/2', '1', '2')


def outgoing(ell, eta, rho, C, sign=1):
    """DLMF H+ (sign=1) or H- (sign=-1) as a complex ball."""
    i = C(0, sign)
    e, r = C(QQ(eta)), C(QQ(rho))
    a, b, z = ell+1+i*e, C(2*ell+2), -2*i*r
    gamma = a.gamma()
    phase = gamma / abs(gamma)
    prefactor = (-i)**ell * (C.pi()*e/2).exp() * phase
    return prefactor * (i*r).exp() * z**(ell+1) * z.hypergeometric_U(a, b)


def enclosed_g(ell, eta, rho, bits):
    C = ComplexBallField(bits)
    g = (outgoing(ell, eta, rho, C, 1) + outgoing(ell, eta, rho, C, -1))/2
    assert g.real().is_finite() and g.imag().is_finite()
    assert g.imag().contains_zero()
    return g.real()


class IrregularCoulomb(numberdb.Generator):
    table = 'T511'
    parameters = ('ell', 'eta', 'rho')
    type = 'R'
    digits = 100
    rigour = 'proven'
    files = ('generate.py', 'checks.py', 'evidence.json', 'README.md')

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

    def value(self, params, digits):
        bits = numberdb.bits(digits + WORKING_EXTRA_DIGITS)
        for attempt in range(4):
            result = enclosed_g(int(params['ell']), params['eta'], params['rho'], bits)
            if result.accuracy() >= numberdb.bits(digits+10):
                return result
            bits *= 2
        raise ArithmeticError('Unable to enclose the requested precision')


if __name__ == '__main__':
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN') == '1':
        os.environ['NUMBERDB_API_KEY'] = sys.stdin.read().strip()
    generator = IrregularCoulomb()
    if '--fill-draft' in sys.argv or os.environ.get('NUMBERDB_PUBLISH') == '1':
        print(generator.publish(message='Add 88 rigorously enclosed Coulomb G values at 100 significant digits'))
    else:
        report = generator.verify(sample=None)
        print(report)
        sys.exit(0 if report.ok else 1)