"""Elliptic dilogarithm companion J at CM torsion points -- numberdb.org/T701.

Run with SageMath or passagemath:
    sage -pip install numberdb
    sage -python generate.py             # verify every stored entry
    sage -python generate.py --submit    # write entries to the existing draft

Submission needs NUMBERDB_API_KEY. It does not offer or publish the draft.
Download checks.py and argument-selection.md for validation and selection.

Brunault (14) fixes J_q. With L=-log|q|, a=r/N, R=|q|, the omitted
one-sided sum starting at exponent t=M+a or M+1-a is bounded by
    L R^t/(1-R^t) * (t/(1-R) + R/(1-R)^2).
Indeed |log|1-z|| <= |z|/(1-|z|) for |z|<1. This follows from the
absolutely convergent logarithm series, and summing (t+k) R^(t+k)
gives the displayed majorant. Both tails are added to the Arb ball.
All arithmetic, including arguments and the Bernoulli term, is enclosed.
The 128 extra working bits allow cancellation and rounding; checks.py
measures the minimum retained accuracy over the entire enumeration.
"""
import os
import sys
from math import gcd
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

DISCRIMINANTS = (-3, -4, -7, -8, -11, -12, -16, -19, -27, -28, -43, -67, -163)
WORKING_GUARD = 128


def parameters():
    for delta in DISCRIMINANTS:
        b = delta % 2
        for order in range(3, 7):
            for r in range(order):
                for s in range(order):
                    if gcd(gcd(order, r), s) != 1:
                        continue
                    if (r, s) > ((-r) % order, (-s) % order):
                        continue
                    if (2*r) % order == 0 and (b*r) % order == 0:
                        continue
                    yield dict(Delta=delta, N=order, r=r, s=s)


def evaluate(delta, order, r, s, digits=100, metadata=False):
    """Evaluate reduced coordinates, including omitted zeros for checks."""
    precision = numberdb.bits(digits, losing=WORKING_GUARD)
    R = RealBallField(precision)
    C = ComplexBallField(precision)
    b = int(delta) % 2
    a = QQ(int(r) % int(order)) / QQ(order)
    s = int(s) % int(order)
    L = R.pi() * R(-int(delta)).sqrt()
    rho = (-L).exp()
    q = (-1 if b else 1) * rho
    phase = (C.gen(0)*R.pi()*QQ(2*s-b*(int(r) % int(order)))/QQ(order)).exp()
    x = C((-L*a).exp()) * phase
    inverse_term = C(q) / x
    value = L**2 * R(a**3 - QQ(3)*a**2/2 + a/2) / 3
    tolerance = R(2)**(-precision+28)

    def bound(t):
        rt = (-L*t).exp()
        return L*rt/(1-rt) * (t/(1-rho)+rho/(1-rho)**2)

    zp, zm = x, inverse_term
    for n in range(1000):
        # The unit-circle factor log|z| is exactly zero, even at z=1.
        if n+a != 0:
            value -= L*(n+a)*(1-zp).log().real()
        value += L*(n+1-a)*(1-zm).log().real()
        M = n+1
        tail = bound(M+a) + bound(M+1-a)
        if tail < tolerance:
            value = value.add_error(tail)
            assert value.is_finite()
            return (value, M, tail) if metadata else value
        zp *= q
        zm *= q
    raise ArithmeticError('Tail did not reach the certified tolerance')


class Table(numberdb.Generator):
    table = 'T701'
    parameters = ('Delta', 'N', 'r', 's')
    type = 'R'
    digits = 100
    rigour = 'proven'

    def enumerate(self):
        yield from parameters()

    def value(self, params, digits):
        values = [int(params[k]) for k in self.parameters]
        # tau^2+tau+1=0 and tau*u=u-1 at u=(tau+2)/3.
        # Reindexing K under lambda -> tau*lambda multiplies K by 1/tau.
        # Periodicity therefore forces K=D=J=0 at this fixed point.
        if values == [-3, 3, 1, 2]:
            return 0
        return evaluate(*values, digits=digits)


if __name__ == '__main__':
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN') == '1':
        numberdb.configure(api_key=sys.stdin.read().strip())
    g = Table()
    if '--submit' in sys.argv or os.environ.get('NUMBERDB_PUBLISH') == '1':
        print(g.publish(message='Compute J in ball arithmetic with explicit tails'))
    else:
        report = g.verify(sample=None)
        print(report)
        sys.exit(0 if report.ok else 1)