"""Power-series coefficients of spheroidal eigenvalues -- numberdb.org/T520.

Run with SageMath:
    sage -pip install numberdb
    sage -python generate.py             # verify every stored entry
    sage -python generate.py --fill      # fill this existing draft with a key

NUMBERDB_API_KEY supplies the writing key. NUMBERDB_KEY_FROM_STDIN=1 reads
it from stdin instead. --fill calls Generator.publish(), which writes entries
and attachments; it does not publish or offer the draft.

The quantity is chi_n(sqrt(q)) = sum_k a_(0,n,k) q^k, q=c^2, for the bounded
order-zero angular problem on [-1,1]. DLMF lambda_n^0(q)=chi_n(sqrt(q))-q.
All arithmetic is rational. The eigenvector gauge is coefficient of P_n = 1;
it does not claim unit L2 norm. Multiplication by x^2 moves a Legendre degree
by at most two, so order r has finite support through degree n+2r. No finite
matrix approximation or numerical differentiation is used.

Selection and proof: argument-selection.md. Independent checks: checks.py.
"""
import numberdb.sage as numberdb
from sage.rings.rational_field import QQ
from functools import lru_cache
import os
import sys

N_MAX = 11
K_MAX = 12


def multiply_x_squared(vector):
    """Apply x^2 in the unnormalised Legendre basis, exactly."""
    result = {}
    for degree, coefficient in vector.items():
        l = degree
        transitions = {
            l + 2: QQ((l + 1) * (l + 2)) / QQ((2*l + 1) * (2*l + 3)),
            l: QQ((l + 1)**2) / QQ((2*l + 1) * (2*l + 3))
               + QQ(l*l) / QQ((2*l + 1) * (2*l - 1)),
        }
        if l >= 2:
            transitions[l - 2] = QQ(l * (l - 1)) / QQ((2*l + 1) * (2*l - 1))
        for target, weight in transitions.items():
            result[target] = result.get(target, QQ(0)) + coefficient * weight
    return {l: v for l, v in result.items() if v}


@lru_cache(maxsize=None)
def perturbation(n, order):
    """Return eigenvalue coefficients and Legendre eigenvector corrections."""
    if n < 0 or order < 0:
        raise ValueError('Mode and Taylor order must be nonnegative')
    eigenvalues = [QQ(n * (n + 1))]
    vectors = [{n: QQ(1)}]
    for r in range(1, order + 1):
        applied = multiply_x_squared(vectors[r - 1])
        eigenvalues.append(applied.get(n, QQ(0)))
        current = {}
        for l in range(n % 2, n + 2*r + 1, 2):
            if l == n:
                continue
            rhs = -applied.get(l, QQ(0))
            for j in range(1, r + 1):
                rhs += eigenvalues[j] * vectors[r - j].get(l, QQ(0))
            value = rhs / QQ(l*(l + 1) - n*(n + 1))
            if value:
                current[l] = value
        vectors.append(current)
    return tuple(eigenvalues), tuple(vectors)


def coefficient(n, k):
    """One exact coefficient, including k=0."""
    return perturbation(int(n), int(k))[0][int(k)]


class SpheroidalCoefficients(numberdb.Generator):
    table = 'T520'
    parameters = ('m', 'n', 'k')
    type = 'Q'
    rigour = 'exact'
    digits = 100  # Irrelevant for exact rationals: no rounding is performed.
    files = ('generate.py', 'checks.py', 'argument-selection.md',
             'numerical-evidence.md', 'skill-manifest.json')

    def enumerate(self, n_max=N_MAX, k_max=K_MAX):
        for n in range(n_max + 1):
            for k in range(k_max + 1):
                yield {'m': 0, 'n': n, 'k': k}

    def value(self, params, digits):
        m, n, k = (int(params[name]) for name in self.parameters)
        if m != 0 or n < 0 or k < 0:
            raise ValueError('This table is restricted to m=0, n>=0, k>=0')
        return perturbation(n, max(K_MAX, k))[0][k]


if __name__ == '__main__':
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN') == '1':
        secret = sys.stdin.read().strip()
        if '=' in secret and secret.split('=', 1)[0].isidentifier():
            secret = secret.split('=', 1)[1].strip().strip('\"').strip("'")
        os.environ['NUMBERDB_API_KEY'] = secret
    generator = SpheroidalCoefficients()
    if '--fill' in sys.argv or os.environ.get('NUMBERDB_FILL') == '1':
        print(generator.publish(message='Exact spheroidal Taylor coefficients, independently checked in two bases'))
    else:
        report = generator.verify(sample=None)
        print(report)
        sys.exit(0 if report.ok else 1)