"""Higher Mahler measures of (x^n-1)/(x-1) -- numberdb.org/T682.

Run with SageMath (or passagemath with the indicated rings):
    sage -pip install numberdb
    sage -python generate.py                   # verify all stored entries
    sage -python checks.py                     # independent offline checks
    sage -python generate.py --fill            # fill an authorized draft

The last command requires NUMBERDB_API_KEY. It sends entries, not a publication
or review action. In the campaign wrapper the key can instead arrive on stdin
with NUMBERDB_KEY_FROM_STDIN=1; NUMBERDB_PUBLISH=1 selects filling.

The geometric-sum roots have rational angles j/n, 1<=j<n. Lalín--Sinha,
arXiv:1106.1304, Theorem 6 gives the third moment from root differences.
Differences r/n occur n-1-r times. The second moment follows from Parseval
for log|sin(n*pi*t)/sin(pi*t)|. No numerical quadrature supplies the bounds.
All transcendental operations use Arb balls, including rational arguments.
WORKING_GUARD=128 adds bits for cancellation; each result is required to retain
32 bits beyond the requested digits. The separate evidence records actual
radii and accuracy. See argument-selection.md for the finite domain decision.
"""
import os
import math
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

WORKING_GUARD = 128
MAX_LENGTH = 50
MOMENT_ORDERS = (2, 3)


@lru_cache(maxsize=None)
def clausen_pair(angle, precision):
    """C_3(2*pi*angle), S_2(2*pi*angle), with exact period/reflection."""
    R = RealBallField(precision)
    C = ComplexBallField(precision)
    a = QQ(angle)
    a -= a.floor()
    if a == 0:
        return R(3).zeta(), R(0)
    if a == QQ(1)/2:
        return -QQ(3)/4 * R(3).zeta(), R(0)
    sign = 1
    if a > QQ(1)/2:
        a = 1-a
        sign = -1
    z = (C(0, 2)*R.pi()*a).exp()
    return z.polylog(3).real(), sign*z.polylog(2).imag()


def moment(n, k, digits=100):
    n, k = int(n), int(k)
    if n < 2 or k not in MOMENT_ORDERS:
        raise ValueError('This implementation requires n>=2 and k=2 or 3.')
    precision = numberdb.bits(digits, losing=WORKING_GUARD)
    R = RealBallField(precision)
    pi = R.pi()
    if k == 2:
        result = pi**2 * QQ(n-1)/QQ(6*n)
    else:
        degree = n-1
        csum, ssum = R(0), R(0)
        for r in range(1, degree):
            angle = QQ(r)/QQ(n)
            c3, s2 = clausen_pair(angle, precision)
            multiplicity = degree-r
            csum += multiplicity*c3
            ssum += multiplicity*s2*(degree*angle-r)
        result = -QQ(3)/2*degree**2*R(3).zeta()-3*degree*csum-3*pi*ssum
    if not result.is_finite() or result.contains_zero():
        raise ArithmeticError('Expected a finite, nonzero enclosure.')
    if not R(result.rad()) < abs(R(result.mid()))*R(2)**(-numberdb.bits(digits)-32):
        raise ArithmeticError('Insufficient retained precision after cancellation.')
    return result


class GeometricSumMahler(numberdb.Generator):
    table = 'T682'
    parameters = ('n', 'k')
    type = 'R'
    digits = 100
    rigour = 'proven'
    files = ('generate.py', 'checks.py', 'readback.py', 'argument-selection.md',
             'numerical-evidence.json', 'skill-manifest.json', 'sources.json')

    def enumerate(self, max_length=MAX_LENGTH):
        for n in range(2, max_length+1):
            for k in MOMENT_ORDERS:
                yield {'n': n, 'k': k}

    def value(self, params, digits):
        n, k = int(params['n']), int(params['k'])
        result = {'number': moment(n, k, digits)}
        if n <= MAX_LENGTH and all(n % divisor for divisor in range(2, math.isqrt(n)+1)):
            companion_n = 1 if n == 2 else n
            result['equals'] = 'HREF{T681#%d,%d}' % (companion_n, k)
        if n == 2:
            result['comment'] = (r'$S_2(x)=1+x$; the moment is $\pi^2/12$.' if k == 2
                                 else r'$S_2(x)=1+x$; the signed third moment is $-3\zeta(3)/2$.')
        return result


if __name__ == '__main__':
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN') == '1':
        numberdb.configure(api_key=sys.stdin.read().strip())
    generator = GeometricSumMahler()
    if '--fill' in sys.argv or os.environ.get('NUMBERDB_PUBLISH') == '1':
        outcome = generator.publish(message='Fill geometric-sum logarithmic moments with certified balls')
        print(outcome)
        print('run:', outcome.run)
    else:
        report = generator.verify(sample=None)
        print(report)
        sys.exit(0 if report.ok else 1)