"""Mahler measures of the Bertin Q_k family -- numberdb.org/T664.

Run with SageMath:
    sage -pip install numberdb mpmath       # install dependencies once
    sage -python generate.py                 # verify every stored entry
    sage -python generate.py --fill          # fill this existing draft only

The default enumeration is integers -8 <= k <= 16, together with k=-36,
ordered by (abs(k), k<0). See argument-selection.md for the separate roles.
Bertin, arXiv:math/0501153, Theorem 1.1(2), supplies the eta quotient and
convergent divisor-sum expansion. The coefficient is -k, and the polynomial
has twelve nonconstant terms. This is the ordinary logarithmic measure.

Solving the eta equation without fixing its branch is insufficient. For
-4<k<0 we choose Re(q)>0, Im(q)<0; for 0<k<12, Re(q)<0, Im(q)<0.
The exterior roots are positive for k<-4 and the small negative root for
k>12. The CM nomes at -36,-6,-3,0,12 are reconstructed exactly from tau.
The endpoint k=-4 uses the proved factorization/Smyth evaluation.

Compute at 150 and 210 decimal digits, independently resolving the nome;
retain 100 significant digits. Agreement and root residuals are numerical
evidence, not a rigorous enclosure for the full computation. An explicit
bound on the omitted q-series is computed separately. No cached values or
rounded parameter constants are used. The generator has no source dependency
beyond numberdb, mpmath, and Python's standard library.
"""
import numberdb.sage as numberdb
import mpmath as mp
import os
import sys
from functools import lru_cache

WORKING_DIGITS = (150, 210)
SELECTED = tuple(sorted(list(range(-8, 17)) + [-36], key=lambda k: (abs(k), k < 0)))


def t_of_nome(q):
    return (q * mp.qp(q**3)**4 * mp.qp(q**12)**8 * mp.qp(q**2)**12
            / (mp.qp(q)**4 * mp.qp(q**4)**8 * mp.qp(q**6)**12))


def nome(k):
    if k == -36:
        return mp.exp(-2 * mp.pi * mp.sqrt(3) / 3)
    if k == -6:
        return mp.exp(-mp.pi * mp.sqrt(3) / 3)
    if k == -3:
        return mp.exp(2 * mp.pi * mp.j * (-3 + mp.j * mp.sqrt(15)) / 24)
    if k == 0:
        return -mp.j * mp.exp(-mp.pi * mp.sqrt(3) / 6)
    if k == 12:
        return -mp.exp(-mp.pi * mp.sqrt(3) / 3)
    a = mp.mpf(-k - 2)
    t = (a - mp.sqrt(a*a - 4)) / 2 if k < 0 else (a + mp.sqrt(a*a - 4)) / 2
    if k < -4:
        seeds = (mp.mpf(1)/10, mp.mpf(1)/5)
    elif k > 12:
        seeds = (-mp.mpf(1)/20, -mp.mpf(1)/10)
    elif k < 0:
        seeds = ((1-mp.j)/4, (mp.mpf(24)/100-mp.j/4))
    else:
        seeds = (-mp.mpf(1)/10-mp.j/5, -mp.mpf(11)/100-mp.j/5)
    q = mp.findroot(lambda q: t_of_nome(q)-t, seeds, maxsteps=100)
    if -4 < k < 0:
        assert q.real > 0 and q.imag < 0
    elif 0 < k < 12:
        assert q.real < 0 and q.imag < 0
    elif k < -4:
        assert 0 < q < 1
    else:
        # This is the root connected to q=0 as k tends to +infinity.
        assert -mp.exp(-mp.pi*mp.sqrt(3)/3) < q < 0
    return q


def q_tail(r, n):
    # sigma_3(j) <= zeta(3)*j^3 < 2*j^3, and the absolute coefficients
    # in the summand total 72. Bound 144*sum_{j>n} j^2*r^j exactly.
    a = n + 1
    return 144*r**a*(a*a/(1-r)+2*a*r/(1-r)**2+r*(1+r)/(1-r)**3)


@lru_cache(None)
def compute(k, working):
    with mp.workdps(working):
        if k == -4:
            val = 7*mp.zeta(3)/mp.pi**2
            return mp.nstr(val, working), {'method':'Smyth factorization'}
        q = nome(k)
        t = t_of_nome(q)
        residual = abs(-t-1/t-2-k)
        assert residual < mp.power(10, -working+8)
        r = abs(q)
        n = int((working+15)*mp.log(10)/(-mp.log(r)))+20
        sig = [0]*(n+1)
        for d in range(1,n+1):
            for j in range(d,n+1,d):
                sig[j] += d**3
        powers = q
        terms=[]
        for j in range(1,n+1):
            terms.append(mp.mpf(sig[j])/j * (-2*powers+16*powers**2+6*powers**3-48*powers**6))
            powers *= q
        val = (-mp.log(q)+mp.fsum(terms)).real
        tail = q_tail(r,n)
        assert tail < mp.power(10,-working-10)
        return mp.nstr(val,working), {
            'q_real':mp.nstr(q.real,working), 'q_imag':mp.nstr(q.imag,working),
            'abs_q':mp.nstr(r,12), 'parameter_residual':mp.nstr(residual,8),
            'series_terms':n, 'q_series_tail_bound':mp.nstr(tail,8)}


COMMENTS = {
    '-4':r'$Q_{-4}=(1+x+xy+xyz)(1+z+zy+zyx)/(xyz)$, so $m(Q_{-4})=7\zeta(3)/\pi^2$ CITE{formula-special}.',
    '-3':r'The Dirichlet evaluation $m(Q_{-3})=\frac85m(1+x+y)$ is proved in CITE{D}, Theorem 1.1.',
    '0':r'The critical parameter $0$ satisfies $m(Q_{12})=4m(Q_0)$ CITE{B}, Theorem 4.1.',
    '12':r'The face-centred cubic spanning-tree entropy per vertex CITE{GR}, Section 5.6.',
    '-6':r'An arithmetic example studied together with $k=-36$ in CITE{B}, Theorem 4.2.',
    '-36':r'An arithmetic example studied together with $k=-6$ in CITE{B}, Theorem 4.2.'}


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

    def enumerate(self):
        for k in SELECTED:
            yield {'k':str(k)}

    def value(self,params,digits):
        if digits > 100:
            raise ValueError('This protocol has been validated for 100 digits.')
        k=int(params['k'])
        if k not in SELECTED:
            raise ValueError('Argument is outside this validated selection.')
        value=numberdb.agreeing(lambda working:compute(k,working)[0],at=WORKING_DIGITS)
        result={'number':value}
        if str(k) in COMMENTS:
            result['comment']=COMMENTS[str(k)]
        return result


if __name__ == '__main__':
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN') == '1':
        numberdb.configure(api_key=sys.stdin.read().strip())
    g=BertinQ()
    if '--fill' in sys.argv or os.environ.get('NUMBERDB_PUBLISH') == '1':
        print(g.publish(message='Fill the twelve-term Bertin Q_k draft after independent checks'))
    else:
        report=g.verify(sample=None)
        print(report)
        sys.exit(0 if report.ok else 1)