"""Values of the Wilton function -- numberdb.org/T678.

Run with SageMath (install numberdb once using sage -pip install numberdb):
    sage -python generate.py             # verify all stored entries
    sage -python generate.py --fill      # write entries with NUMBERDB_API_KEY

The API method called publish writes entries to this existing draft; it does
not publish or offer the table. Only its owner controls that later lifecycle.
The single parameter is the exact positive root of A*t^2+B*t-K, named
quad_A_B_mK. See argument-selection.md for the arithmetic selection and bounds.
All arithmetic determining orbits is exact integer arithmetic. Arb encloses
radicals, logarithms, and the signed finite-cycle quotient. WORKING_GUARD=96
provides extra working bits; checks.py measures the achieved accuracy and
independently encloses the defining series with a certified absolute tail.
No external data or helper modules control the selection.
"""
import sys
from math import gcd, isqrt
import numberdb.sage as numberdb
from sage.rings.real_arb import RealBallField

WORKING_GUARD = 96
DISCRIMINANT_BOUND = 100


def discriminant(t):
    A, B, K = t
    return B*B + 4*A*K


def arguments(bound=DISCRIMINANT_BOUND):
    out = []
    for A in range(1, isqrt(bound)+1):
        for B in range(1, isqrt(bound)+1):
            for K in range(1, (bound-1)//4+1):
                D = B*B + 4*A*K
                if (D <= bound and isqrt(D)**2 != D
                        and gcd(gcd(A, B), K) == 1
                        and A+B-K > 0 and A-B-K < 0):
                    out.append((A, B, K))
    return sorted(out, key=lambda t: (discriminant(t), t))


def identity(t):
    return 'quad_%d_%d_m%d' % t


def decode(key):
    prefix, A, B, K = key.split('_')
    if prefix != 'quad' or not K.startswith('m'):
        raise ValueError('Expected quad_A_B_mK')
    t = (int(A), int(B), int(K[1:]))
    A, B, K = t
    D = discriminant(t)
    if not (A > 0 and B > 0 and K > 0 and gcd(gcd(A, B), K) == 1
            and isqrt(D)**2 != D and A+B-K > 0 and A-B-K < 0):
        raise ValueError('Argument is not a primitive purely periodic quadratic')
    return t


def step(t):
    A, B, K = t
    m = (B+isqrt(discriminant(t)))//(2*K)
    return m, (K, 2*m*K-B, A+B*m-K*m*m)


def orbit(t):
    seen, rows, digits = set(), [], []
    u = t
    while u not in seen:
        seen.add(u)
        rows.append(u)
        a, u = step(u)
        digits.append(a)
    if u != t:
        raise ValueError('Unexpected preperiod')
    return rows, digits


def root(t, F):
    A, B, K = t
    return F(2*K)/(F(B)+F(discriminant(t)).sqrt())


def cycle_value(t, digits=100):
    F = RealBallField(numberdb.bits(digits, losing=WORKING_GUARD))
    rows, partials = orbit(t)
    product, total = F(1), F(0)
    for j, u in enumerate(rows):
        x = root(u, F)
        total += (-1)**j * product * (-x.log())
        product *= x
    result = total/(1-(-1)**len(rows)*product)
    if not result.is_finite():
        raise ArithmeticError('Nonfinite enclosure')
    return result


def label(t):
    A, B, K = t
    return r'$\frac{-' + str(B) + r'+\sqrt{' + str(discriminant(t)) + '}}{' + str(2*A) + '}$'


class Wilton(numberdb.Generator):
    table = 'T678'
    parameters = ('x',)
    type = 'R'
    digits = 100
    rigour = 'proven'

    def enumerate(self, bound=DISCRIMINANT_BOUND):
        for t in arguments(bound):
            yield {'x': identity(t)}

    def value(self, params, digits):
        t = decode(params['x'])
        result = {'number': cycle_value(t, digits), 'param-latex': label(t)}
        if t == (1, 1, 1):
            result['comment'] = r'The golden-section argument is $x=1/\varphi=[0;\overline1]$, and $W(x)=\log(\varphi)/\varphi$.'
        elif t == (1, 2, 1):
            result['comment'] = r'The argument is $x=\sqrt2-1=[0;\overline2]$, and $W(x)=\log(1+\sqrt2)/\sqrt2$.'
        return result


if __name__ == '__main__':
    g = Wilton()
    if '--fill' in sys.argv:
        print(g.publish(message='Compute 141 Wilton values with certified cycle and series enclosures'))
    else:
        report = g.verify(sample=None)
        print(report)
        sys.exit(0 if report.ok else 1)