"""Prolate spheroidal differential eigenvalues -- numberdb.org/T518.

Run with SageMath and the NumberDB client:
    sage -pip install numberdb
    sage -python generate.py                  # verify all stored entries
    sage -python generate.py --fill           # fill an authorized draft
    sage -python checks.py                    # independent numerical checks

The fill operation writes entries and attached source files; it does not
publish or offer a draft. Set NUMBERDB_API_KEY for authenticated operations.
The campaign wrapper can instead supply the key on stdin with
NUMBERDB_KEY_FROM_STDIN=1 and request filling with NUMBERDB_PUBLISH=1.

For m=0 the self-adjoint operator is -d/dx((1-x*x)d/dx) + c*c*x*x.
Its matrix in normalized Legendre polynomials separates by parity. We use
Sturm bisection on each finite Jacobi matrix, comparing working precision
and truncation independently. This is heuristic agreement, not a certified
infinite-operator bound. checks.py solves the endpoint ODE independently.
"""
import os
import sys
from functools import lru_cache
from pathlib import Path
import numberdb.sage as numberdb
from sage.rings.rational_field import QQ
from sage.rings.real_mpfr import RealField

BANDWIDTHS = (1, 2, 4, 8)
MODES = 12
SETTINGS = ((160, 80), (220, 80), (220, 120), (220, 180))
REQUIRED_AGREEMENT = 110


def field(dps):
    return RealField(numberdb.bits(dps, losing=16))


@lru_cache(maxsize=None)
def rational_matrix(parity, size):
    diagonal, off_squared = [], []
    for j in range(size):
        k = 2*j + parity
        x2 = QQ((k+1)**2) / ((2*k+1)*(2*k+3))
        x2 += QQ(k*k) / ((2*k-1)*(2*k+1))
        diagonal.append((k*(k+1), x2))
        if j < size-1:
            off_squared.append(QQ((k+1)**2 * (k+2)**2) /
                               ((2*k+1)*(2*k+3)**2*(2*k+5)))
    return tuple(diagonal), tuple(off_squared)


def matrix(q, parity, size):
    R = q.parent()
    diagonal, off_squared = rational_matrix(parity, size)
    return ([R(k) + q*R(a) for k, a in diagonal],
            [q*q*R(b) for b in off_squared])


def sturm_count(diagonal, off_squared, x):
    """Number of eigenvalues strictly below x (zero pivots by a limit)."""
    pivot = diagonal[0] - x
    count = int(pivot < 0)
    tiny = x.parent()(2)**(-x.parent().precision()*2)
    if not pivot:
        pivot = -tiny
        count = 1
    for i, square in enumerate(off_squared):
        pivot = diagonal[i+1] - x - square/pivot
        if not pivot:
            pivot = -tiny
        count += int(pivot < 0)
    return count


def bisect_mode(q, n, size, dps):
    R = q.parent()
    diagonal, off_squared = matrix(q, n % 2, size)
    index = n // 2
    lo = R(n*(n+1))-1
    hi = R(n*(n+1))+q+1
    assert sturm_count(diagonal, off_squared, lo) <= index
    assert sturm_count(diagonal, off_squared, hi) > index
    tolerance = R(10)**(-(dps-10))
    while hi-lo > tolerance:
        mid = (lo+hi)/2
        if sturm_count(diagonal, off_squared, mid) <= index:
            lo = mid
        else:
            hi = mid
    return (lo+hi)/2


@lru_cache(maxsize=None)
def spectrum(S, dps, size):
    R = field(dps)
    c = R.pi()*R(QQ(S))/2
    return tuple(bisect_mode(c*c, n, size, dps) for n in range(MODES))


def decimal_string(x, digits):
    # Fixed significant digits, retaining a decimal point, including final zeros.
    return x.str(truncate=False, digits=digits, skip_zeroes=False)


class ProlateDifferentialEigenvalues(numberdb.Generator):
    table = 'T518'
    parameters = ('m', 'S', 'n')
    type = 'R'
    digits = 100
    rigour = 'heuristic (agreement-checked)'
    files = ('generate.py', 'checks.py', 'evidence.json')

    def enumerate(self):
        for S in BANDWIDTHS:
            for n in range(MODES):
                yield {'m': '0', 'S': str(QQ(S)), 'n': str(n)}

    def value(self, params, digits):
        if int(params['m']) != 0:
            raise ValueError('Only angular order zero is supported')
        if digits != 100:
            raise ValueError('Higher precision requires new measured settings')
        S, n = str(params['S']), int(params['n'])
        values = [spectrum(S, dps, size)[n] for dps, size in SETTINGS]
        R = field(220)
        reference = R(values[-1])
        for other in values[:-1]:
            assert abs(R(other)-reference) < abs(reference)*R(10)**(-REQUIRED_AGREEMENT)
        return decimal_string(reference, digits)


if __name__ == '__main__':
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN') == '1':
        os.environ['NUMBERDB_API_KEY'] = sys.stdin.read().strip()
    generator = ProlateDifferentialEigenvalues()
    if '--fill' in sys.argv or os.environ.get('NUMBERDB_PUBLISH') == '1':
        # Compute all entries before sending so they arrive as one batch.
        for params in generator.enumerate():
            generator.value(params, generator.digits)
        print(generator.publish(message='48 differential eigenvalues, checked by precision, truncation and endpoint ODE'))
    else:
        report = generator.verify(sample=None)
        print(report)
        sys.exit(0 if report.ok else 1)
