"""Independent checks for Neumann eigenvalues -- numberdb.org/T485.

    sage -python checks.py             # compute and check locally
    sage -python checks.py --stored    # read draft; NUMBERDB_API_KEY required

The external triangle mode list is Table 2, p. 393, of Laugesen, Pan and Son,
Mathematical Inequalities & Applications 15 (2012), 381-394,
doi:10.7153/mia-15-32. The integers are transcribed independently of generate.py.
Bessel mode selection here uses mpmath, independently of scipy in generate.py.
Stored decimal endpoints are checked with the separate identity
x J'_m(x) = m J_m(x) - x J_(m+1)(x). No generator root/bracket helper is used.
A quadrature-based cosine Rayleigh quotient checks rectangle and right
triangle values, starting with a nonzero known-answer control.
"""
import json
import math
import os
import sys
from pathlib import Path
from fractions import Fraction
from functools import lru_cache

import numberdb.sage as numberdb
from numberdb._write import to_text
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
import scipy
from generate import NeumannEigenvalues, DOMAINS, COUNT

R = RealBallField(640)
C = ComplexBallField(640)
I = RealIntervalField(640)
EQ = [(0,0),(0,1),(1,0),(1,1),(0,2),(2,0),(1,2),(2,1),(0,3),(3,0),
      (2,2),(1,3),(3,1),(0,4),(4,0),(2,3),(3,2),(1,4),(4,1),(0,5)]
HE = [(0,0),(0,1),(1,1),(0,2),(1,2),(0,3),(2,2),(1,3),(0,4),(2,3),
      (1,4),(0,5),(3,3),(2,4),(1,5),(0,6),(3,4),(2,5),(1,6),(4,4)]


def rat(text):
    x = Fraction(str(text))
    return QQ(x.numerator)/QQ(x.denominator)


def endpoints(text):
    """Exact interpretation of the site's decimal ± one last-place unit."""
    from decimal import Decimal
    if '.' not in text and 'e' not in text.lower():
        return rat(text), rat(text)
    x = Decimal(text)
    unit = QQ(10)**x.as_tuple().exponent
    q = rat(text)
    return q-unit, q+unit


def finite_narrow(ball, digits=98):
    assert ball.is_finite(), ball
    assert ball == 0 or ball.accuracy() >= math.floor(digits*math.log2(10)), ball


def contains_decimal(text, expected):
    a,b = endpoints(text)
    # The comparator must itself carry at least 120 digits, not be a huge ball.
    finite_narrow(expected, 120)
    if a == b:
        assert expected == R(a), (text, expected)
    else:
        assert expected > R(a) and expected < R(b), (text, expected)


def source_triangle(qpairs):
    return [R(QQ(16*(j*j+j*k+k*k))/9)*R.pi()**2 for j,k in qpairs]


@lru_cache(None)
def rayleigh_cos(k, length):
    """Compute a one-dimensional eigenvalue from energy integrals."""
    if k == 0:
        return mp.mpf(0)
    wave = k*mp.pi/length
    norm = mp.quad(lambda x: mp.cos(wave*x)**2,[0,length])
    energy = mp.quad(lambda x: (wave*mp.sin(wave*x))**2,[0,length])
    assert norm > 0
    return energy/norm


def quadratic_spectra():
    # Increase the candidate eigenvalue coefficient and solve one square
    # equation. This is separate from the generator's two-index sorting.
    out={d:[] for d in ('unit-square','rectangle-1x2','right-isosceles-triangle')}
    modes={d:[] for d in out}
    for q in range(1000):
        for d in out:
            if len(out[d]) >= COUNT:
                continue
            scale=4 if d=='rectangle-1x2' else 1
            hits=[]
            for j in range(math.isqrt(q//scale)+1):
                k=math.isqrt(q-scale*j*j)
                if scale*j*j+k*k != q:
                    continue
                if d=='right-isosceles-triangle' and j<k:
                    continue
                hits.append((j,k))
            for j,k in hits:
                if len(out[d]) == COUNT:
                    break
                out[d].append(R(QQ(q)/scale)*R.pi()**2)
                modes[d].append((j,k))
        if all(len(v)==COUNT for v in out.values()):
            return out,modes
    raise AssertionError('enumeration incomplete')


def independent_bessel_spectra():
    # A cutoff beyond every retained eigenvalue; all angular orders >= 30
    # have first root >30 (DLMF 10.21.3). Each computed branch includes
    # its first mode above the cutoff, which checks radial exhaustion.
    cutoff=mp.mpf(600)
    modes=[]
    for m in range(25):
        k=1
        while True:
            x=mp.besseljzero(m,k+(m==0),derivative=1)
            if x*x>cutoff:
                break
            modes.append((x*x,m,k))
            k+=1
    out={}
    for d,step in [('unit-disk',1),('half-disk',1),('quarter-disk',2),('sector-60',3)]:
        candidates=[(mp.mpf(0),0,0)]
        for v,m,k in modes:
            if m%step==0:
                candidates.extend([(v,m,k)]*(2 if d=='unit-disk' and m>0 else 1))
        candidates.sort()
        assert len(candidates)>COUNT and candidates[COUNT][0]<cutoff
        out[d]=candidates[:COUNT]
    return out


def bracket_stored(text,m):
    a,b=endpoints(text)
    assert a>0 and b>a
    def g(t):
        z=C(R(t).sqrt())
        return (m*z.bessel_J(m)-z*z.bessel_J(m+1)).real()
    ga,gb=g(a),g(b)
    assert ga.is_finite() and gb.is_finite()
    assert (ga>0 and gb<0) or (ga<0 and gb>0),(m,text,ga,gb)


def main(stored=False):
    mp.mp.dps=140
    # Known nonzero energy; ensures quadrature has not coerced everything to 0.
    control=rayleigh_cos(1,1)
    assert abs(control-mp.pi**2)<mp.mpf('1e-130')
    assert abs(control-mp.mpf('9.86960440108935861883449'))<mp.mpf('1e-23')
    print('Known-answer control: cosine energy quotient = pi^2; PASS',flush=True)
    if stored:
        doc=numberdb.table('T485')
        entries={(d,n):v for d,rows in doc['Numbers'].items() for n,v in rows.items()}
        values={key:(v['number'] if isinstance(v,dict) else v) for key,v in entries.items()}
    else:
        g=NeumannEigenvalues()
        entries={(p['domain'],str(p['n'])):g.value(p,100) for p in g.enumerate()}
        values={key:to_text(v['number'],100) for key,v in entries.items()}
        print('Minimum nonzero ball accuracy:', min(v['number'].accuracy() for v in entries.values() if v['number'] != 0), 'bits',flush=True)
    assert set(values)=={(d,str(n)) for d in DOMAINS for n in range(1,COUNT+1)}
    assert all(values[d,'1']=='0' for d in DOMAINS)
    print('180 identities and nine exact zero modes; PASS',flush=True)

    spectra,modes=quadratic_spectra()
    spectra['equilateral-triangle']=source_triangle(EQ)
    spectra['hemiequilateral-triangle']=source_triangle(HE)
    for d,rows in spectra.items():
        for n,v in enumerate(rows,1):
            contains_decimal(values[d,str(n)],v)
    print('100 polygonal entries, including 40 against published triangle modes; PASS',flush=True)

    for d,rows in modes.items():
        length=2 if d=='rectangle-1x2' else 1
        for n,(j,k) in enumerate(rows,1):
            v=rayleigh_cos(j,1)+rayleigh_cos(k,length)
            expected=R(rat(mp.nstr(v,135)))
            contains_decimal(values[d,str(n)],expected)
            # Reflection-invariant cosine modes for the right triangle:
            # x+y=1 uses a parity factor (-1)^(j+k) in the reflected term.
            if d=='right-isosceles-triangle':
                def u(x,y):
                    return (mp.cos(j*mp.pi*x)*mp.cos(k*mp.pi*y)
                            +(-1)**(j+k)*mp.cos(k*mp.pi*x)*mp.cos(j*mp.pi*y))
                for t in (mp.mpf(1)/7,mp.mpf(2)/5):
                    assert abs(mp.diff(lambda x:u(x,t),0))<mp.mpf('1e-125')
                    assert abs(mp.diff(lambda y:u(t,y),0))<mp.mpf('1e-125')
                    normal=mp.diff(lambda x:u(x,1-t),t)+mp.diff(lambda y:u(t,y),1-t)
                    assert abs(normal)<mp.mpf('1e-125')
    print('60 energy-quotient checks and all right-triangle boundary checks; PASS',flush=True)

    with mp.workdps(40):
        circular=independent_bessel_spectra()
    root_ids=set()
    for d,rows in circular.items():
        for n,(v,m,k) in enumerate(rows,1):
            if k==0: continue
            text=values[d,str(n)]
            with mp.workdps(40):
                assert abs(mp.mpf(text)-v)<mp.mpf('1e-32')
            bracket_stored(text,m)
            root_ids.add((m,k))
            if d!='unit-disk':
                alpha={'half-disk':mp.pi,'quarter-disk':mp.pi/2,'sector-60':mp.pi/3}[d]
                assert abs(m*mp.sin(m*alpha))<mp.mpf('1e-125')
    print('76 positive circular entries: independent ordering, multiplicities, and written-endpoint brackets; PASS',flush=True)

    t22=numberdb.table('T22')['Numbers']
    comparisons=0
    for m,k in sorted(root_ids):
        index=str(k+(m==0))
        row=t22.get(str(m),{}).get(index)
        if row is None: continue
        text=row.get('number') if isinstance(row,dict) else row
        a,b=endpoints(text)
        # Independent corpus values, squared without shortening their bounds.
        lo,hi=a*a,b*b
        for d,rows in circular.items():
            for n,(_,mm,kk) in enumerate(rows,1):
                if (mm,kk)==(m,k):
                    aa,bb=endpoints(values[d,str(n)])
                    assert max(aa,lo)<=min(bb,hi)
                    comparisons+=1
    assert comparisons == sum(m <= 10 and k > 0 for rows in circular.values() for _,m,k in rows), comparisons
    print(str(comparisons)+' circular entries within T22 coverage agree with its squared entries and zero-index convention; PASS',flush=True)

    for d in DOMAINS:
        rows=[endpoints(values[d,str(n)]) for n in range(1,COUNT+1)]
        assert all(a[0]<=b[1] for a,b in zip(rows,rows[1:]))
        assert rows[1][0]>0
    print('Ordering and multiplicity-preserving layout; PASS',flush=True)
    t390=numberdb.table('T390')['Numbers']
    equalities=0
    import re
    for key,entry in entries.items():
        if not isinstance(entry,dict) or not entry.get('equals'): continue
        target=re.fullmatch(r'HREF\{([^#]*)#([^,]+),(\d+)\}',entry['equals'])
        assert target, entry
        table,domain,index=target.groups()
        if table:
            assert table=='Dirichlet_eigenvalues_of_the_classical_planar_domains'
            v=t390[domain][index]
            other=v['number'] if isinstance(v,dict) else v
        else:
            other=values[domain,index]
        a,b=endpoints(values[key]);c,d=endpoints(other)
        assert max(a,c)<=min(b,d),(key,entry['equals'])
        equalities+=1
    print(str(equalities)+' exact-identity links agree with their read-back targets; PASS',flush=True)

    print(json.dumps({'mode':'stored' if stored else 'local','entries':len(values),
       'distinct_positive_bessel_roots':len(root_ids),'T22_comparisons':comparisons,'equality_links':equalities,
       'longest_written_entry':max(map(len,values.values())),
       'written_characters':sum(map(len,values.values())),
       'numberdb':numberdb.__version__,'mpmath':mp.__version__,'scipy':scipy.__version__}))
    print('NUMERICAL_VALUES_JSON '+json.dumps({d:{str(n):values[d,str(n)] for n in range(1,COUNT+1)} for d in DOMAINS}))


if __name__=='__main__':
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN')=='1':
        os.environ['NUMBERDB_API_KEY']=sys.stdin.read().strip()
    main('--stored' in sys.argv or os.environ.get('NUMBERDB_CHECK_STORED')=='1')