"""Moments of uniform random walks in the plane -- numberdb.org/T522.

Run with SageMath (or passagemath with numberdb):
    sage -pip install numberdb
    sage -python generate.py             # verify every stored entry
    sage -python generate.py --publish   # fill this draft using NUMBERDB_API_KEY

The publish method writes entries; it does not offer or publish the draft.
The default selection is n=3,4 and s=-1,1,3,5,7; see argument-selection.md.
All arithmetic is in Arb balls. Three-step seeds use proved gamma formulas.
Four-step seeds use elliptic integrals split at k=1/2. On (0,1/2),
DLMF 19.12.1 gives a convergent logarithmic series, integrated termwise;
the square-series remainder has an explicit positive bound. On [1/2,1],
Arb quadrature encloses the analytic integrand, with a checked branch domain.
See numerical-method.md for the error proof and references. Higher odd
moments use the proved recurrences in Borwein et al. (2011), Example 1.
"""
import os
import sys
from functools import lru_cache
import numberdb.sage as numberdb
from sage.rings.rational_field import QQ
from sage.rings.real_arb import RealBallField
from sage.rings.complex_arb import ComplexBallField


def finite_narrow(x, digits):
    if not x.is_finite() or not x > 0:
        raise ArithmeticError('Nonfinite or nonpositive moment enclosure')
    if not x.rad() < abs(x.mid()) * (QQ(1) / 10**(digits + 8)):
        raise ArithmeticError('Moment enclosure is too wide')
    return x


def log_integrals(R, q):
    """Integrals of k^q * (1, L, L^2) over (0,1/2), L=log(4/k)."""
    b = R(q + 1)
    h = R(8).log()
    t = R(2)**(-q - 1) / b
    return t, t * (h + 1/b), t * (h*h + 2*h/b + 2/(b*b))


def elliptic_square_integrals(bits):
    """Enclose J_p=int_0^1 k^p K(1-k^2)^2 dk for p=0,2 (parameter K)."""
    R = RealBallField(bits)
    C = ComplexBallField(bits)
    terms = bits // 2 + 12
    # c_j = ((1/2)_j/j!)^2; d_j = 2(H_{2j}-H_j).
    c, d = [R(1)], [R(0)]
    for j in range(1, terms):
        c.append(c[-1] * (QQ(2*j-1) / QQ(2*j))**2)
        d.append(d[-1] + QQ(1) / QQ(j*(2*j-1)))
    # Square sum c_j*k^(2j)*(L-d_j) and integrate exactly in balls.
    A = [R(0) for _ in range(2*terms-1)]
    B = [R(0) for _ in A]
    D = [R(0) for _ in A]
    for j in range(terms):
        for l in range(j, terms):
            factor = c[j]*c[l] * (1 if j == l else 2)
            A[j+l] += factor
            B[j+l] += factor*(d[j]+d[l])
            D[j+l] += factor*d[j]*d[l]
    result = []
    for p in (0, 2):
        low = R(0)
        for r in range(len(A)):
            i0, i1, i2 = log_integrals(R, 2*r+p)
            low += A[r]*i2 - B[r]*i1 + D[r]*i0
        # 0 <= K^2-S^2 <= (32/9)*L^2*k^(2*terms), for k <= 1/2.
        error = (R(32)/9)*log_integrals(R, 2*terms+p)[2]
        low += R(0, error.upper())
        def integrand(k, analytic):
            # K(1-k^2) is analytic for Re(k)>0. Reject complex boxes
            # that reach Re(k)<=0, even when the integrand is bounded.
            if analytic and not k.real() > 0:
                return C('nan')
            return k**p * (1-k*k).elliptic_k()**2
        high = C.integral(integrand, QQ(1)/2, 1,
                          abs_tol=R(2)**(-bits+24),
                          rel_tol=R(2)**(-bits+24),
                          eval_limit=100000, depth_limit=40)
        if not high.real().is_finite() or not high.imag().is_finite() or not high.imag().contains_zero():
            raise ArithmeticError('Elliptic quadrature failed')
        result.append(low + high.real())
    return tuple(result)


@lru_cache(maxsize=None)
def moments(n, working_digits=150):
    bits = numberdb.bits(working_digits, losing=48)
    R = RealBallField(bits)
    pi = R.pi()
    if n == 3:
        a = R(3)/16 * R(2)**(QQ(1)/3) * R(QQ(1)/3).gamma()**6 / pi**4
        b = a + R(27)/4 * R(2)**(QQ(2)/3) * R(QQ(2)/3).gamma()**6 / pi**4
    elif n == 4:
        j0, j2 = elliptic_square_integrals(bits)
        a = 4*j0/pi**3
        b = 16*(j0-3*j2)/pi**3
    else:
        raise ValueError('The implemented seed formulas require n=3 or n=4')
    values = {-1:a, 1:b}
    for s in (-1,1,3):
        if n == 3:
            v = (2*(5*s*s+30*s+46)*values[s+2]-9*(s+2)**2*values[s]) / R((s+4)**2)
        else:
            v = (4*(s+3)*(5*s*s+30*s+48)*values[s+2]-64*(s+2)**3*values[s]) / R((s+4)**3)
        values[s+4] = v
    for value in values.values():
        finite_narrow(value, working_digits-10)
    return values


class PlanarWalkMoments(numberdb.Generator):
    table = 'T522'
    parameters = ('n', 's')
    type = 'R'
    digits = 100
    rigour = 'proven'
    files = ('generate.py', 'checks.py', 'check.py', 'argument-selection.md',
             'numerical-method.md', 'reference-values.json',
             'instructions-read.json', 'skill-manifest.json')

    def enumerate(self):
        for n in (3, 4):
            # Mean distance first; the inverse-distance moment follows it.
            for s in (1, -1, 3, 5, 7):
                yield {'n': n, 's': s}

    def value(self, params, digits):
        n, s = int(params['n']), int(params['s'])
        if s not in (-1,1,3,5,7):
            raise ValueError('The implemented range is s=-1,1,3,5,7')
        value = finite_narrow(moments(n, max(150, digits+50))[s], digits)
        if s == 1:
            return {'number':value, 'comment':'Mean distance from the origin.'}
        if s == -1:
            return {'number':value, 'comment':'Mean reciprocal distance from the origin.'}
        return value


if __name__ == '__main__':
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN') == '1':
        key = sys.stdin.read().strip()
        if '=' in key and key.split('=',1)[0].isidentifier():
            key = key.split('=',1)[1].strip().strip('\"\'')
        numberdb.configure(api_key=key)
    generator = PlanarWalkMoments()
    if '--publish' in sys.argv or os.environ.get('NUMBERDB_PUBLISH') == '1':
        print(generator.publish(message='Fill ten certified radial moments from unit-step planar walks.'))
    else:
        report = generator.verify(sample=None)
        print(report)
        sys.exit(0 if report.ok else 1)