"""Higher Mahler measures of 1+x_1+...+x_(n-1) -- numberdb.org/T686.

Run with SageMath:
    sage -pip install numberdb mpmath        # once
    sage -python generate.py                 # verify all stored entries
    sage -python generate.py --fill          # fill the existing private draft

The generator computes raw signed logarithmic moments, with normalized Haar
measure, natural logarithms, and no factorial divisor. The default selection
is n=3,4 and k=2,...,8. See argument-selection.md and method.md.
All arithmetic in truncated series uses Arb. Series tails are not certified:
150- and 220-decimal working computations are compared, and the declared
rigour is heuristic (agreement-checked), not proven.
"""
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.polynomial.polynomial_ring_constructor import PolynomialRing
from math import factorial

WORKING_DIGITS = (150, 220)
MAX_ORDER = 8


def falling(x, j):
    z = QQ(1)
    for i in range(j):
        z *= x-i
    return z


def ode_coefficients(center, direction):
    P = PolynomialRing(QQ, 'u'); u = P.gen(); x = center+direction*u
    return [list(x**4-64), list(direction*(7*x**5-32*x**3+64*x)),
            list(6*x**6-60*x**4), list(direction*(x**7-20*x**5+64*x**3))]


def frobenius(a, exponent, seeds, count, R):
    """Coefficient recurrence from BSWZ (11), expanded about 2 or 4."""
    c = [R(v) for v in seeds]
    r = QQ(exponent)
    for m in range(len(c), count):
        denom = R(0); rhs = R(0)
        for j, poly in enumerate(a):
            for ell, v in enumerate(poly):
                if not v:
                    continue
                i = m+j-ell-2
                if i < 0 or i > m:
                    continue
                z = v*falling(r+i, j)
                if i == m:
                    denom += z
                else:
                    rhs += z*c[i]
        assert not denom.contains_zero()
        c.append(-rhs/denom)
    return c


def low4(count, R):
    """BSWZ (26),(30); p4(x)=sum (r_j-s_j log(x))*x^(2j+1)."""
    s = [3/(2*R.pi()**2)]
    r = [s[0]*R(8).log()]
    for j in range(1, count):
        a=4*(2*j-1)*(5*j*j-5*j+2); b=-2*(j-1)**3
        sm2=s[j-2] if j>1 else R(0)
        rm2=r[j-2] if j>1 else R(0)
        sj=(a*s[-1]+b*sm2)/(128*j**3)
        rj=(a*r[-1]+b*rm2+3*(64*j*j*sj-(20*j*j-20*j+6)*s[-1]+(j-1)**2*sm2))/(128*j**3)
        s.append(sj);r.append(rj)
    return r,s


def density4_series(working):
    R=RealBallField(numberdb.bits(3*working+50, losing=96))
    P=PolynomialRing(R,'u'); N=4*working
    a=ode_coefficients(2,1)
    A=P(frobenius(a,0,[1,0],N,R))
    B=P(frobenius(a,0,[0,1],N,R))
    S=P(frobenius(a,QQ(1)/2,[1],N,R))*(-4/R.pi()**2)
    p2=R(2)**(QQ(7)/3)*R.pi()/(3*R(3).sqrt()*R(QQ(2)/3).gamma()**6)
    p1=R(5).sqrt()/(40*R.pi()**4)
    for j in (1,2,4,8):p1*=R(QQ(j)/15).gamma()
    slope=(p1-p2*A(-1))/B(-1)
    C=p2*A+slope*B
    F=P(frobenius(ode_coefficients(4,-1),QQ(1)/2,[R(2).sqrt()/R.pi()**2],N,R))
    r,s=low4(2*working,R)
    return R,P,N,r,s,C,S,F


def moments4(working, max_order=MAX_ORDER):
    R,P,N,r,s,C,S,F=density4_series(working)
    log2=P([R(2).log()]+[R((-1)**(j+1))/(j*2**j) for j in range(1,N)])
    log4=P([R(4).log()]+[-R(1)/(j*4**j) for j in range(1,N)])
    pow2=P(1);pow4=P(1);values=[]
    for k in range(max_order+1):
        low=sum((r[j]*((-1)**k*factorial(k))/R(2*j+2)**(k+1)
                 -s[j]*((-1)**(k+1)*factorial(k+1))/R(2*j+2)**(k+2)) for j in range(len(r)))
        cp=(C*pow2).truncate(N);sp=(S*pow2).truncate(N);fp=(F*pow4).truncate(N)
        mid=sum(2*cp[j]/(j+1) for j in range(0,N,2))
        high=sum((sp[j]+fp[j])/(R(j)+QQ(3)/2) for j in range(N))
        values.append(low+mid+high)
        pow2=(pow2*log2).truncate(N);pow4=(pow4*log4).truncate(N)
    return values


def moments3(working,max_order=MAX_ORDER):
    """Taylor coefficients of BBSW Theorem 4.1, equation (23)."""
    R=RealBallField(numberdb.bits(working,losing=96));P=PolynomialRing(R,'s');s=P.gen()
    N=max_order+1;term=P(1);total=term
    for j in range(2*working):
        # (1+s/2+j)^3 / (4*(j+1)^2*(j+3/2+s/2)).
        b=R(j)+QQ(3)/2
        inverse=P([(-R(1)/(2*b))**k/b for k in range(N)])
        term=(term*(j+1+s/2)**3*inverse).truncate(N)/(4*(j+1)**2)
        total+=term
    # log(3^(s+1) Gamma(1+s/2)^2/Gamma(2+s)).
    logarithm=[R(3).log(),R(3).log()-1]
    for k in range(2,N):
        logarithm.append(((-1)**k*(R(2)**(1-k)-1)*R(k).zeta()+(-1)**k)/k)
    exponential=[R(3)]
    for k in range(1,N):
        exponential.append(sum(j*logarithm[j]*exponential[k-j] for j in range(1,k+1))/k)
    result=(P(exponential)*total).truncate(N)*R(3).sqrt()/(2*R.pi())
    return [result[k]*factorial(k) for k in range(N)]


@lru_cache(None)
def computed(n, working):
    v=(moments3 if n==3 else moments4)(working)
    # These are decimal approximations to truncated computations. Their Arb
    # radii bound arithmetic only, not the omitted series tails.
    import mpmath as mp
    with mp.workdps(working+10):
        return tuple(mp.nstr(mp.mpf(x.mid().str(truncate=False)),working) for x in v)


class HigherMahler(numberdb.Generator):
    table='T686'
    parameters=('n','k')
    type='R'
    digits=100
    rigour='heuristic (agreement-checked)'

    def enumerate(self):
        for n in (3,4):
            for k in range(2,9):
                yield {'n':n,'k':k}

    def value(self, params, digits):
        n,k=int(params['n']),int(params['k'])
        if n not in (3,4) or not 2<=k<=8:
            raise ValueError('Implemented selection: n=3,4 and 2<=k<=8.')
        return numberdb.agreeing(lambda w:computed(n,w)[k],at=WORKING_DIGITS)


if __name__=='__main__':
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN')=='1':
        numberdb.configure(api_key=sys.stdin.read().strip())
    generator=HigherMahler()
    if '--fill' in sys.argv or os.environ.get('NUMBERDB_PUBLISH')=='1':
        print(generator.publish(message='Computed 14 higher Mahler measures with independent controls'))
    else:
        report=generator.verify(sample=None);print(report)
        sys.exit(0 if report.ok else 1)
