"""Neumann eigenvalues of the classical planar domains -- numberdb.org/T485.

Run with SageMath (or passagemath with scipy and mpmath):
    sage -pip install numberdb scipy mpmath
    sage -python generate.py                 # verify all stored entries
    sage -python generate.py --fill          # fill this draft, key in environment
    sage -python checks.py                   # independent local checks

The first 20 eigenvalues include the constant mode and all multiplicities.
Integer quadratic forms enumerate the integrable polygons. scipy's double
precision Bessel derivative roots select low disk and sector modes; every
selected root is recomputed with mpmath at 130 and 160 decimal digits, then
bracketed by exact rational endpoints using Arb evaluations of J'_m.
The endpoint bracket, rather than the mpmath point estimate, is squared.

The precision guard is 192 bits. The exact root bracket has radius
10^(-digits-12), leaving over ten guard digits after squaring. This is measured
by checks.py. Index identification is agreement checked, not interval certified;
the declared overall rigour is therefore heuristic (agreement-checked).
"""
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.real_arb import RealBallField
from sage.rings.complex_arb import ComplexBallField
from sage.rings.real_mpfi import RealIntervalField
import mpmath as mp
from scipy.special import jnp_zeros

TABLE = 'T485'
COUNT = 20
WORKING_GUARD = 192
ORDER_BOUND = 40
RADIAL_BOUND = 20
DOMAINS = (
    'unit-square', 'rectangle-1x2', 'unit-disk', 'half-disk',
    'quarter-disk', 'sector-60', 'equilateral-triangle',
    'right-isosceles-triangle', 'hemiequilateral-triangle',
)
CIRCULAR = {'unit-disk': 1, 'half-disk': 1, 'quarter-disk': 2, 'sector-60': 3}


def rational_decimal(text):
    """Convert decimal text to an exact rational without float coercions."""
    from fractions import Fraction
    x = Fraction(str(text))
    return QQ(x.numerator) / QQ(x.denominator)


def derivative(order, x, C):
    z = C(x)
    if order == 0:
        return -z.bessel_J(1).real()
    return (z.bessel_J(order - 1) - z.bessel_J(order + 1)).real() / 2


@lru_cache(None)
def root_enclosure(order, k, digits):
    # mpmath counts the zero at the origin for order zero; scipy does not.
    index = k + (order == 0)
    with mp.workdps(digits + 30):
        low_precision = mp.besseljzero(order, index, derivative=1)
    with mp.workdps(digits + 60):
        high_precision = mp.besseljzero(order, index, derivative=1)
        assert abs(low_precision - high_precision) < mp.mpf(10)**(-digits - 20)
        center = rational_decimal(mp.nstr(high_precision, digits + 20))
    radius = QQ(1) / (ZZ(10)**(digits + 12))
    a, b = center - radius, center + radius
    bits = numberdb.bits(digits, losing=WORKING_GUARD)
    R, C, I = RealBallField(bits), ComplexBallField(bits), RealIntervalField(bits)
    fa, fb = derivative(order, R(a), C), derivative(order, R(b), C)
    assert fa.is_finite() and fb.is_finite()
    assert (fa > 0 and fb < 0) or (fa < 0 and fb > 0), (order, k)
    root = R(I(a, b))
    assert root.is_finite() and root > 0 and root.accuracy() > numberdb.bits(digits + 9, losing=0)
    return root


@lru_cache(None)
def layouts():
    """Mode labels; integer/rational sorting for all polygonal domains."""
    result = {}
    for domain in DOMAINS:
        if domain in CIRCULAR:
            continue
        rows = []
        for j in range(COUNT + 1):
            for k in range(COUNT + 1):
                if domain in ('right-isosceles-triangle', 'hemiequilateral-triangle') and k > j:
                    continue
                if domain in ('equilateral-triangle', 'hemiequilateral-triangle'):
                    q = QQ(16) * (j*j + j*k + k*k) / 9
                elif domain == 'rectangle-1x2':
                    q = QQ(j*j) + QQ(k*k) / 4
                else:
                    q = QQ(j*j + k*k)
                rows.append((q, j, k))
        rows.sort()
        # Any omitted j or k is >= COUNT+1; even the smallest coefficient
        # (1/4 in the rectangle) then exceeds the last retained coefficient.
        assert rows[COUNT - 1][0] < QQ((COUNT + 1)**2) / 4
        result[domain] = [('pi', q, j, k) for q, j, k in rows[:COUNT]]

    roots = {m: jnp_zeros(m, RADIAL_BOUND) for m in range(ORDER_BOUND + 1)}
    for domain, step in CIRCULAR.items():
        candidates = [(0.0, 0, 0, 0)]
        for m in range(0, ORDER_BOUND + 1, step):
            multiplicity = 2 if domain == 'unit-disk' and m > 0 else 1
            for k, root in enumerate(roots[m], 1):
                candidates.extend((float(root)**2, m, k, copy) for copy in range(multiplicity))
        candidates.sort()
        cutoff = candidates[COUNT - 1][0]
        # DLMF 10.21.3: the first positive derivative root for m>0 is >m.
        assert cutoff < (ORDER_BOUND + 1)**2
        assert all(float(roots[m][-1])**2 > cutoff + 1 for m in range(0, ORDER_BOUND + 1, step))
        result[domain] = [('zero',)] + [('bessel', m, k) for _, m, k, _ in candidates[1:COUNT]]
    return result


class NeumannEigenvalues(numberdb.Generator):
    table = TABLE
    parameters = ('domain', 'n')
    type = 'R'
    digits = 100
    rigour = 'heuristic (agreement-checked)'
    files = ('generate.py', 'checks.py')

    def enumerate(self):
        for domain in DOMAINS:
            for n in range(1, COUNT + 1):
                yield {'domain': domain, 'n': n}

    def number(self, params, digits):
        mode = layouts()[params['domain']][int(params['n']) - 1]
        if mode[0] == 'zero' or (mode[0] == 'pi' and mode[1] == 0):
            return ZZ(0)
        if mode[0] == 'pi':
            R = RealBallField(numberdb.bits(digits, losing=WORKING_GUARD))
            return R(mode[1]) * R.pi()**2
        return root_enclosure(mode[1], mode[2], digits)**2

    def value(self, params, digits):
        result = {'number': self.number(params, digits)}
        domain, n = params['domain'], int(params['n'])
        mode = layouts()[domain][n - 1]
        # These exact mode identities give the corresponding Dirichlet
        # fundamental eigenvalue. They are also checked against stored T390.
        fundamental = {
            'unit-square': QQ(2), 'rectangle-1x2': QQ(5)/4,
            'equilateral-triangle': QQ(16)/3,
            'right-isosceles-triangle': QQ(5),
            'hemiequilateral-triangle': QQ(112)/9,
        }
        if mode[0] == 'pi' and mode[1] == fundamental.get(domain):
            result['equals'] = 'HREF{Dirichlet_eigenvalues_of_the_classical_planar_domains#%s,1}' % domain
        elif mode == ('bessel', 0, 1):
            result['equals'] = 'HREF{Dirichlet_eigenvalues_of_the_classical_planar_domains#half-disk,1}'
        else:
            # Repeat multiplicities and shared exact modal values without
            # mistaking numerical closeness for equality.
            def identity(item):
                if item[0] == 'zero' or (item[0] == 'pi' and item[1] == 0):
                    return ('zero',)
                return item[:2] if item[0] == 'pi' else item
            wanted = identity(mode)
            for other in DOMAINS:
                for i, candidate in enumerate(layouts()[other], 1):
                    if (other, i) == (domain, n):
                        return result
                    if identity(candidate) == wanted:
                        result['equals'] = 'HREF{#%s,%s}' % (other, i)
                        return result
        return result


if __name__ == '__main__':
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN') == '1':
        os.environ['NUMBERDB_API_KEY'] = sys.stdin.read().strip()
    generator = NeumannEigenvalues()
    if '--fill' in sys.argv or os.environ.get('NUMBERDB_PUBLISH') == '1':
        print(generator.publish(message='Fill the checked first 20 Neumann eigenvalues of nine domains.'))
    else:
        report = generator.verify(sample=None)
        print(report)
        sys.exit(0 if report.ok else 1)
