"""Normalized second moments of the lattice quantizers -- numberdb.org/T498.

Run with SageMath (install numberdb once with sage -pip install numberdb):
    sage -python generate.py                  # verify every stored entry
    sage -python checks.py                    # independent geometric checks
    sage -python generate.py --fill-draft     # fill T498 with NUMBERDB_API_KEY

The table holds G = integral_V |x|^2 dx / (n vol(V)^(1+2/n)).
Conway and Sloane, IEEE Trans. Inform. Theory 28 (1982), 211-226,
DOI 10.1109/TIT.1982.1056483, equations (32)-(36), (60)-(61), supplies
A_n, D_n, E_n, and A_n*. For D_n* we integrate its clipped-cube Voronoi
cell exactly by inclusion-exclusion (their Section III-F describes this
cell). Root lattices have minimum norm 2, but G is invariant under scale.
The infinite families run through dimension 24, the common endpoint of the
classical-lattice packing and covering tables. Exceptional dual, laminated,
and Coxeter-Todd lattices are outside this initial computed range.

Every coefficient is rational. Only a positive rational root is evaluated
in Arb, at numberdb.bits(digits, losing=64); the check script measures the
relative radii and checks the full range at 100 and 150 significant digits.
Rational outputs (Z^n, A_1, E_8) are returned exactly, not as decimals.
"""
import os
import sys
from functools import lru_cache
from math import comb, factorial
import numberdb.sage as numberdb
from sage.rings.rational_field import QQ
from sage.rings.real_arb import RealBallField

WORKING_GUARD = 64
MAX_DIMENSION = 24


@lru_cache(None)
def permutohedron_j(n):
    """Conway-Sloane (60), with J_n = I(P_{n-1})/n."""
    if n == 1:
        return QQ(0)
    first = QQ(factorial(n)) / (8 * (n + 1) * n ** (n - 2))
    first *= sum(QQ(n ** k) / factorial(k) for k in range(n + 1))
    first -= QQ(n * n) / (4 * (n + 1))
    rest = sum(QQ(comb(n, r)) * QQ((r, n)) ** r
               * QQ((n - r, n)) ** (n - r - 1) * permutohedron_j(r)
               for r in range(1, n))
    return first + QQ((1, n + 1)) * rest


@lru_cache(None)
def dual_d_integral(n):
    """Integral of norm squared over |x_i|<=1/2, sum |x_i|<=n/4.

    In each orthant let y_i=2|x_i|. On the unit cube cut by sum y_i<=n/2,
    shift j coordinates by 1 in inclusion-exclusion. Dirichlet simplex
    integrals give 2n*s^(n+2)/(n+2)! + 2j*s^(n+1)/(n+1)! + j*s^n/n!.
    The orthant multiplicity cancels the coordinate Jacobian; norm^2
    contributes 1/4. The Voronoi cell has volume 1/2.
    """
    total = QQ(0)
    for j in range((n - 1) // 2 + 1):
        s = QQ((n, 2)) - j
        total += (-1) ** j * comb(n, j) * (
            QQ((2 * n, factorial(n + 2))) * s ** (n + 2)
            + QQ((2 * j, factorial(n + 1))) * s ** (n + 1)
            + QQ((j, factorial(n))) * s ** n)
    return total / 4


def exact_form(family, n):
    """Return rational a,b with G = a * b^(1/n), before rounding."""
    if family == 'Z':
        return QQ((1, 12)), QQ(1)
    if family == 'A':
        return QQ((1, 12)) + QQ((1, 6 * (n + 1))), QQ((1, n + 1))
    if family == 'A*':
        return permutohedron_j(n + 1) / (n * (n + 1)), QQ(n + 1)
    if family == 'D':
        return QQ((1, 12)) + QQ((1, 2 * n * (n + 1))), QQ((1, 4))
    if family == 'D*':
        return 2 * dual_d_integral(n) / n, QQ(4)
    if family == 'E':
        return {6: (QQ((5, 56)), QQ((1, 3))),
                7: (QQ((163, 2016)), QQ((1, 2))),
                8: (QQ((929, 12960)), QQ(1))}[n]
    raise ValueError('No formula for this lattice')


def evaluate(family, n, digits=100):
    a, b = exact_form(family, n)
    # These are exactly the rational cases in the selected domain.
    if b == 1:
        return a
    if n == 1:
        return a * b
    field = RealBallField(numberdb.bits(digits, losing=WORKING_GUARD))
    return field(a) * field(b) ** QQ((1, n))


class LatticeSecondMoments(numberdb.Generator):
    table = 'T498'
    parameters = ('family', 'n')
    type = 'R'
    digits = 100
    rigour = 'proven'
    files = ('generate.py', 'checks.py', 'check.py')

    def enumerate(self):
        for family, lower in [('Z', 1), ('A', 1), ('A*', 3), ('D', 4), ('D*', 5)]:
            for n in range(lower, MAX_DIMENSION + 1):
                yield {'family': family, 'n': n}
        for n in (6, 7, 8):
            yield {'family': 'E', 'n': n}

    def value(self, params, digits):
        f, n = params['family'], int(params['n'])
        value = evaluate(f, n, digits)
        comment = {('Z', 1): r'The interval has $G(\mathbb{Z})=1/12$.',
                   ('Z', 2): r'The square lattice.',
                   ('Z', 3): r'The simple cubic lattice.',
                   ('A', 1): r'This one-dimensional lattice is similar to $\mathbb{Z}$.',
                   ('A', 2): r'The hexagonal lattice, with $G(A_2)=5/(36\sqrt{3})$.',
                   ('A', 3): r'The face-centered cubic lattice.',
                   ('A*', 3): r'The body-centered cubic lattice, with $G(A_3^*)=19/(192\sqrt[3]{2})$.',
                   ('D', 4): r'The Voronoi cell is the regular $24$-cell.',
                   ('E', 8): r'The exact value is $G(E_8)=929/12960$.'}.get((f,n))
        return {'number': value, 'comment': comment} if comment else value


if __name__ == '__main__':
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN') == '1':
        os.environ['NUMBERDB_API_KEY'] = sys.stdin.read().strip()
    generator = LatticeSecondMoments()
    if '--fill-draft' in sys.argv or os.environ.get('NUMBERDB_PUBLISH') == '1':
        print(generator.publish(message='Fill the draft with certified lattice second moments'))
    else:
        report = generator.verify(sample=None)
        print(report)
        sys.exit(0 if report.ok else 1)