"""Asymptotic amplitudes of the classes of trees -- numberdb.org/T482.

Run with SageMath (or passagemath):
    sage -pip install numberdb
    sage -python generate.py                 # verify the stored draft/table
    sage -python generate.py --compute       # compute without network writes
    sage -python generate.py --fill          # fill this draft, API key required

The --fill operation writes entries, not publication or review status.
Set NUMBERDB_API_KEY for access to the unpublished draft.

All fourteen classes match the node-counting conventions of T448. In
particular, a planted root is included in n. HRS (1975), corrected (1986),
and Genitrini (2016, arXiv:1605.00837) give the singularity method.

Certified arithmetic: a_n <= Catalan(n-1) <= 4^(n-1) for each rooted
unordered class, by forgetting the order of plane trees. At y < 1/4,
truncating A(y) after M gives tail <= y(4y)^M/(1-4y); truncating yA'(y)
gives tail <= y(4y)^M*((M+1)-4My)/(1-4y)^2. Every omitted tail is added as
a symmetric Arb error ball. For Q=sum_{k>=2} sign_k A(x^k)/k, the tails
after K are bounded by x^(K+1)/((K+1)(1-x)(1-4x^(K+1))) and, for xQ',
by x^(K+1)/((1-x)(1-4x^(K+1))^2). The constants 1/2 and 23/100 bound x
and x^k respectively in the summation cutoff choices. Exact inequalities
are verified at runtime; these choices are not empirical error estimates.

Newton steps only propose a rational centre. Opposite certified signs at
centre +/- 10^(-digits-15) certify the root interval. The asymptotic formulas
are evaluated on that interval, retaining the root error. Working precision
has 160 guard bits, and series errors target digits+35 decimal places.
The build also repeated the computation with 30 extra requested digits.
"""
import os
import sys
from functools import lru_cache

if os.environ.get('NUMBERDB_KEY_FROM_STDIN') == '1':
    os.environ['NUMBERDB_API_KEY'] = sys.stdin.read().strip()

import numberdb.sage as numberdb
from sage.rings.integer_ring import ZZ
from sage.rings.rational_field import QQ
from sage.rings.real_arb import RealBallField

CLASSES = ('rooted', 'free', 'planted', 'identity-rooted', 'identity-free',
           'series-reduced-planted', 'series-reduced-rooted',
           'series-reduced-free', 'degree-bounded-2', 'degree-bounded-3',
           'plane', 'plane-binary', 'labelled-rooted', 'labelled-free')
GUARD_BITS = 160


@lru_cache(None)
def coefficients(kind, N):
    """Exact integer coefficients; each division has its remainder checked."""
    a = [ZZ(0)] * (N + 1)
    if kind in ('rooted', 'identity', 'series'):
        sigma = [ZZ(0)] * (N + 1)
        for n in range(1, N + 1):
            if n == 1:
                a[n] = ZZ(1)
            else:
                total = sum((sigma[j] + ((-1)**j if kind == 'series' else 0))
                            * a[n-j] for j in range(1, n))
                a[n], remainder = total.quo_rem(n-1)
                assert remainder == 0 and a[n] >= 0
            for j in range(n, N + 1, n):
                sign = (-1)**(j//n + 1) if kind == 'identity' else 1
                sigma[j] += sign * n * a[n]
    else:
        degree = int(kind)
        square = [ZZ(0)] * (N+1)
        for n in range(1, N+1):
            m = n-1
            square[m] = sum(a[i]*a[m-i] for i in range(1,m))
            total = 6*(1 if m == 0 else a[m])
            total += 3*(square[m] + (a[m//2] if m%2 == 0 else 0))
            if degree == 3:
                total += sum(square[i]*a[m-i] for i in range(1,m))
                total += 3*sum(a[j]*a[m-2*j] for j in range(1,m//2+1))
                total += 2*(a[m//3] if m%3 == 0 else 0)
            a[n], remainder = ZZ(total).quo_rem(6)
            assert remainder == 0 and a[n] >= 0
    return tuple(a)


class Enclosure:
    def __init__(self, kind, digits):
        self.kind = kind
        self.B = RealBallField(numberdb.bits(digits, losing=GUARD_BITS))
        self.eps = self.B(QQ(1)/10**(digits+35))
        self.N = 1
        upper = {'rooted':QQ(34)/100, 'identity':QQ(41)/100,
                 'series':QQ(46)/100, '2':QQ(41)/100, '3':QQ(36)/100}[kind]
        self.ymax = upper**2
        q = self.B(4*self.ymax)
        while not self.B(self.ymax)*q**self.N*(self.N+1)/(1-q)**2 < self.eps:
            self.N += 1
        self.a = coefficients(kind, self.N)
        self.K = 4*(digits+35)
        self.cutoffs = {}

    def series(self, y, k):
        B = self.B
        assert y > 0 and y < self.ymax
        # The cutoff is a speed choice only: the actual tail is bounded anew.
        M = self.cutoffs.get(k)
        if M is None:
            M = 1
            q = 4*y
            while not y*q**M*(M+1)/(1-q)**2 < self.eps:
                M += 1
            self.cutoffs[k] = M
        assert M <= self.N
        value = B(0)
        weighted = B(0)
        for n in range(M, 0, -1):
            value = (value + self.a[n])*y
            weighted = (weighted + n*self.a[n])*y
        q = 4*y
        tail = y*q**M/(1-q)
        dtail = y*q**M*((M+1)-M*q)/(1-q)**2
        return value + B(0,tail.upper()), weighted + B(0,dtail.upper())

    def Q(self, x, absolute=False):
        B = self.B
        value, weighted = B(0), B(0)
        y = x
        for k in range(2,self.K+1):
            y *= x
            A, D = self.series(y,k)
            sign = (-1)**(k+1) if self.kind == 'identity' and not absolute else 1
            value += sign*A/k
            weighted += sign*D
        y *= x
        tail = y/((self.K+1)*(1-x)*(1-4*y))
        dtail = y/((1-x)*(1-4*y)**2)
        return value+B(0,tail.upper()), weighted+B(0,dtail.upper())

    def equation(self, x):
        """Return g(x), g'(x), and the squared square-root coefficient."""
        if self.kind in ('rooted','identity','series'):
            Q, D = self.Q(x)
            g = x.log()+1+Q
            b2half = 1+D
            if self.kind == 'series':
                g -= (1+x).log()
                b2half -= x/(1+x)
            return g, b2half/x, 2*b2half
        A2, D2 = self.series(x*x,2)
        if self.kind == '2':
            tau = 1/x-1
            g = (1-x)**2/x**2-2-A2
            derivative = -2*(1-x)/x**3-2*D2/x
            b2 = 2*(tau/x + D2)
            return g, derivative, b2
        A3, D3 = self.series(x**3,3)
        tau = -1+(2/x-1-A2).sqrt()
        phi = 1+tau+(tau**2+A2)/2+(tau**3+3*tau*A2+2*A3)/6
        g = x*phi-tau
        Fx = phi+(1+tau)*D2+D3
        return g, Fx, 2*Fx/(1+tau)

    def solve(self, digits):
        seeds = {'rooted':QQ(338)/1000, 'identity':QQ(397)/1000,
                 'series':QQ(457)/1000,'2':QQ(403)/1000,'3':QQ(355)/1000}
        if self.kind == 'identity':
            # |x Q'(x)| is bounded by its positive-term majorant at 0.4.
            # Consequently g is strictly increasing on (0,0.4].
            _, majorant = self.Q(self.B(QQ(2)/5), absolute=True)
            assert majorant < 1
        x = self.B(seeds[self.kind])
        for _ in range(8):
            g, d, _ = self.equation(x)
            x = self.B((x-g/d).mid())
        centre = QQ(x.mid())
        radius = QQ(1)/10**(digits+15)
        low, high = self.B(centre-radius), self.B(centre+radius)
        gl, _, _ = self.equation(low)
        gh, _, _ = self.equation(high)
        assert (gl < 0 < gh) or (gh < 0 < gl), (self.kind,gl,gh)
        rho = self.B(centre,radius)
        _, derivative, b2 = self.equation(rho)
        assert not derivative.contains_zero() and b2 > 0
        b = b2.sqrt()
        assert rho.is_finite() and b.is_finite()
        return rho, b


@lru_cache(None)
def compute(digits=100):
    B = RealBallField(numberdb.bits(digits, losing=GUARD_BITS))
    sqrtpi = B.pi().sqrt()
    out, radii, evidence = {}, {}, {}
    for kind in ('rooted','identity','series','2','3'):
        engine = Enclosure(kind,digits)
        rho, b = engine.solve(digits)
        c = b/(2*sqrtpi)
        radii[kind] = rho
        evidence[kind] = {'N':engine.N,'max_used':max(engine.cutoffs.values()),
                          'K':engine.K,'bits_retained':int(c.accuracy())}
        if kind == 'rooted':
            out.update({'rooted':c,'free':b**3/(4*sqrtpi),'planted':rho*c})
        elif kind == 'identity':
            out.update({'identity-rooted':c,'identity-free':b**3/(4*sqrtpi)})
        elif kind == 'series':
            out.update({'series-reduced-planted':rho*c,'series-reduced-rooted':c,
                        'series-reduced-free':(1+rho)*b**3/(4*sqrtpi)})
        else:
            out['degree-bounded-'+kind] = c
    out['plane'] = 1/(4*sqrtpi)
    out['plane-binary'] = B(3).sqrt()/(2*sqrtpi)
    out['labelled-rooted'] = 1/(2*B.pi()).sqrt()
    out['labelled-free'] = out['labelled-rooted']
    for v in out.values():
        assert v.is_finite() and v > 0 and v.accuracy() > numberdb.bits(digits)
    return out, radii, evidence


COMMENTS = {
 'rooted':r'$\alpha=3/2$. Unlabelled rooted trees, with unordered children; this is Otter\'s rooted-tree amplitude CITE{RootedAmplitude}.',
 'free':r'$\alpha=5/2$. Unlabelled unrooted trees; this is Otter\'s unrooted-tree amplitude CITE{FreeAmplitude}.',
 'planted':r'$\alpha=3/2$. The planted root is an additional node, so the count on $n$ nodes equals the rooted-tree count on $n-1$ nodes.',
 'identity-rooted':r'$\alpha=3/2$. Only the identity automorphism preserves the root. The underlying unrooted tree need not be asymmetric CITE{Genitrini}.',
 'identity-free':r'$\alpha=5/2$. The unrooted tree has trivial automorphism group CITE{Correction}.',
 'series-reduced-planted':r'$\alpha=3/2$. The planted root has degree $1$ and is counted as a node; no other node has degree $2$ CITE{HRS}.',
 'series-reduced-rooted':r'$\alpha=3/2$. No node has degree $2$, including the root CITE{SeriesRooted}.',
 'series-reduced-free':r'$\alpha=5/2$. Unrooted trees with no node of degree $2$ CITE{HRS}.',
 'degree-bounded-2':r'$\alpha=3/2$. Every node has at most two unordered children. The count on $n$ nodes is A001190$(n+1)$ CITE{Binary}.',
 'degree-bounded-3':r'$\alpha=3/2$. Every node has at most three unordered children CITE{Ternary}.',
 'plane':r'$\alpha=3/2$. Children are ordered. The count on $n$ nodes is the Catalan number $C_{n-1}$ CITE{Catalan}.',
 'plane-binary':r'$\alpha=3/2$. Each node has zero, one or two ordered children. The count on $n$ nodes is the Motzkin number $M_{n-1}$ CITE{Motzkin}.',
 'labelled-rooted':r'$\alpha=3/2$. The node labels are $1,\ldots,n$, and Cayley\'s count is $n^{n-1}$ CITE{Cayley}.',
 'labelled-free':r'$\alpha=5/2$. The node labels are $1,\ldots,n$, and Cayley\'s count is $n^{n-2}$ CITE{Cayley}.',
}
# Apostrophes in raw single-quoted strings must not reach the document escaped.
COMMENTS = {k:v.replace("\\'", "'") for k,v in COMMENTS.items()}


class TreeAmplitudes(numberdb.Generator):
    table = 'T482'
    parameters = ('class','size','labelling')
    type = 'R'
    digits = 100
    rigour = 'proven'
    files = ('generate.py', 'checks.py', 'check.py', 'dry_run.py', 'build-report.md')

    def enumerate(self):
        for name in CLASSES:
            yield {'class':name,'size':'nodes',
                   'labelling':'labelled' if name.startswith('labelled-') else 'unlabelled'}

    def value(self, params, digits):
        return {'number':compute(digits)[0][params['class']],
                'comment':COMMENTS[params['class']]}


if __name__ == '__main__':
    if '--compute' in sys.argv:
        values, radii, evidence = compute()
        for name in CLASSES:
            print(name, values[name])
        print('evidence',evidence)
    elif '--fill' in sys.argv or os.environ.get('NUMBERDB_PUBLISH') == '1':
        print(TreeAmplitudes().publish(message='Fill fourteen tree amplitudes with certified 100-digit enclosures'))
    else:
        result = TreeAmplitudes().verify(sample=None)
        print(result)
        sys.exit(0 if result.ok else 1)
