"""Laguerre–Selberg integrals for classical Wishart ensembles -- numberdb.org/T564.

Run with SageMath or passagemath and the numberdb client:
    sage -pip install numberdb
    sage -python generate.py          # verify all stored entries
    sage -python checks.py           # independent moment checks
    sage -python generate.py --fill   # fill an authorized draft, API key in env

The full labelled positive orthant is integrated with weight exp(-x), without
an ordered-chamber factor. Gamma factors reduce to q * sqrt(pi)**e exactly.
Only nonzero e uses real ball arithmetic. The 96 guard bits are measured by
checks.py; no precision is inferred from numerical zeros. Selection and bounds
are documented in argument-selection.md. --fill writes entries, never offers,
reviews or publishes a draft. NUMBERDB_KEY_FROM_STDIN=1 reads an API key from
stdin; NUMBERDB_PUBLISH=1 is the wrapper's equivalent of --fill.
"""
import os
import sys
from math import factorial
import numberdb.sage as numberdb
from sage.rings.integer_ring import ZZ
from sage.rings.rational_field import QQ
from sage.rings.real_arb import RealBallField

WORKING_GUARD = 96


def selected_shapes(small_n=6, small_r=6, near_n=16, aspect_n=10):
    shapes = {(n, r) for n in range(1, small_n + 1)
              for r in range(small_r + 1)}
    shapes.update((n, r) for n in range(1, near_n + 1) for r in range(3))
    shapes.update((n, k*n) for n in range(1, aspect_n + 1) for k in (1, 2))
    return sorted(shapes)


def gamma_reduction(twice_argument):
    """Return (rational coefficient, sqrt(pi) exponent) for positive k/2."""
    k = int(twice_argument)
    if k <= 0:
        raise ValueError('Gamma argument must be positive')
    if k % 2 == 0:
        return QQ(factorial(k//2 - 1)), 0
    h = (k - 1)//2
    return QQ(factorial(2*h)) / (ZZ(4)**h * ZZ(factorial(h))), 1


def reduced_integral(beta, n, r):
    """Exact reduction of DLMF 5.14.5 with insertion index zero."""
    beta, n, r = int(beta), int(n), int(r)
    if beta not in (1, 2, 4) or n < 1 or r < 0:
        raise ValueError('Require beta in {1,2,4}, n >= 1, r >= 0')
    q, e = QQ(1), 0
    den_q, den_e = gamma_reduction(2 + beta)
    for j in range(n):
        for k in (beta*(r + 1 + j), 2 + beta*(j + 1)):
            coeff, power = gamma_reduction(k)
            q *= coeff
            e += power
        q /= den_q
        e -= den_e
    return q, e


class LaguerreSelberg(numberdb.Generator):
    table = 'T564'
    parameters = ('beta', 'n', 'r')
    type = 'R'
    digits = 100
    rigour = 'proven'
    files = ('generate.py', 'checks.py', 'argument-selection.md',
             'numerical-evidence.json', 'skill-evidence.json', 'check.py')

    def enumerate(self):
        for beta in (1, 2, 4):
            for n, r in selected_shapes():
                yield {'beta': str(beta), 'n': str(n), 'r': str(r)}

    def value(self, params, digits):
        q, e = reduced_integral(*(int(params[k]) for k in self.parameters))
        if e == 0:
            return q
        field = RealBallField(numberdb.bits(digits, losing=WORKING_GUARD))
        return field(q) * field.pi().sqrt()**e


if __name__ == '__main__':
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN') == '1':
        os.environ['NUMBERDB_API_KEY'] = sys.stdin.read().strip()
    generator = LaguerreSelberg()
    if '--fill' in sys.argv or os.environ.get('NUMBERDB_PUBLISH') == '1':
        print(generator.publish(message='Fill verified Laguerre–Selberg integral draft'))
    else:
        report = generator.verify(sample=None)
        print(report)
        sys.exit(0 if report.ok else 1)