"""Low-degree growth constants of finite p-group algebras -- numberdb.org/T654.

With SageMath and numberdb installed:
    sage -python generate.py                 # verify all stored entries
    sage -python generate.py --fill          # write numbers to this draft
    sage -python checks.py                  # independent numerical checks

The invariant is min H(t)/t^(degree(H)/3) on 0<t<=1, over the prime
coefficient field. See argument-selection.md for every selection decision.
Exact integer arithmetic brackets the sole stationary root. Arb then encloses
the objective on that entire bracket. 80 extra working bits cover bracket
propagation and polynomial evaluation; checks.py measures the retained accuracy.
The one-third threshold is fixed, not an enumerated parameter.
"""
import os
import sys
from functools import lru_cache
import numberdb.sage as numberdb
from sage.rings.integer_ring import ZZ
from sage.rings.rational_field import QQ
from sage.rings.polynomial.polynomial_ring_constructor import PolynomialRing
from sage.rings.real_arb import RealBallField

R = PolynomialRing(ZZ, 't')
t = R.gen()
WORKING_GUARD = 80


def selected_groups():
    for p in (2, 3, 5, 7):
        for n in range(3, 7):
            if (p-1)*n*(n*n-1)//6 <= 80:
                yield 'UT-%d-%d' % (n, p)
        for r in range(2, 5):
            yield 'Heis-%d-%d' % (r, p)


def group_data(identity):
    family, size, prime = identity.split('-')
    size, p = int(size), ZZ(prime)
    if not p.is_prime():
        raise ValueError('The coefficient field must be a prime field')
    if family == 'UT' and size >= 2:
        degrees = [(j, size-j) for j in range(1, size)]
        label = r'$U_{%d}(\mathbb F_{%d})$' % (size, p)
    elif family == 'Heis' and size >= 1:
        degrees = [(1, 2*size), (2, 1)]
        label = r'$\operatorname{Heis}_{%d}(\mathbb F_{%d})$' % (size, p)
    else:
        raise ValueError('Expected UT-n-p (n>=2) or Heis-r-p (r>=1)')
    return p, degrees, label


@lru_cache(None)
def hilbert(identity):
    p, degrees, _ = group_data(identity)
    h = R.one()
    for weight, multiplicity in degrees:
        h *= sum((t**(weight*k) for k in range(p)), R.zero())**multiplicity
    return h


def homogeneous_sign(coefficients, numerator, exponent):
    """Sign at numerator/2^exponent, with no rounding or rational divisions."""
    total = ZZ(coefficients[-1])
    for i, coefficient in enumerate(reversed(coefficients[:-1]), 1):
        total = total*numerator + (ZZ(coefficient) << (exponent*i))
    return total.sign()


def root_bracket(h, steps):
    """Enclose the unique root of 3*t*H'(t)-deg(H)*H(t) in (0,1)."""
    d = h.degree()
    q = [ZZ((3*j-d)*h[j]) for j in range(d+1)]
    assert q[0] < 0 and sum(q) > 0
    # There is exactly one sign change, so Descartes' rule proves uniqueness.
    signs = [c.sign() for c in q if c]
    assert sum(a != b for a,b in zip(signs,signs[1:])) == 1
    low = ZZ.zero()
    for k in range(1, steps+1):
        mid = 2*low+1
        sign = homogeneous_sign(q, mid, k)
        if sign == 0:
            root = QQ(mid)/ZZ(2)**k
            return root, root
        low = mid if sign < 0 else mid-1
    den = ZZ(2)**steps
    return QQ(low)/den, QQ(low+1)/den


def enclose(h, digits=100):
    precision = numberdb.bits(digits, losing=WORKING_GUARD)
    a, b = root_bracket(h, precision)
    ball = RealBallField(precision)
    midpoint, radius = (a+b)/ZZ(2), (b-a)/ZZ(2)
    root = ball(midpoint).add_error(ball(radius))
    assert root > 0 and root < 1
    value = h(root) * (-(ball(h.degree())/3)*root.log()).exp()
    assert value.is_finite() and value > 1 and value < h(1)
    return value, (a, b), root


class GrowthConstants(numberdb.Generator):
    table = 'T654'
    parameters = ('G',)
    type = 'R'
    digits = 100
    rigour = 'proven'
    files = ('generate.py', 'checks.py', 'argument-selection.md',
             'derivation.md', 'skill-manifest.json', 'numerical-evidence.json',
             'companion-hilbert.json', 'check.py')

    def enumerate(self):
        for identity in selected_groups():
            yield {'G': identity}

    def value(self, params, digits):
        identity = params['G']
        value, _, _ = enclose(hilbert(identity), digits)
        return {'number':value, 'param-latex':group_data(identity)[2]}


if __name__ == '__main__':
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN') == '1':
        numberdb.configure(api_key=sys.stdin.read().strip())
    generator = GrowthConstants()
    if '--fill' in sys.argv or os.environ.get('NUMBERDB_FILL') == '1':
        print(generator.publish(message='Compute 25 growth constants with certified stationary-root brackets'))
    else:
        report = generator.verify(sample=None)
        print(report)
        sys.exit(0 if report.ok else 1)