"""Independent checks for Slepian concentration eigenvalues -- numberdb.org/T519.

    sage -python checks.py
    sage -python checks.py --stored  # API key needed while T519 is a draft

Gauss-Legendre Nystrom eigenvalues use the integral kernel directly; no
Legendre differential matrix, eigenvectors or Fourier-ratio formula is used.
Known-answer controls precede comparisons. JSON evidence is printed to stdout.
"""
import json
import os
import sys
import time
import platform
import numberdb.sage as numberdb
import mpmath as mp
from generate import spectrum, agreement_digits, rounded, SETTINGS, SHANNON_NUMBERS, MODES


def gauss_nodes(ctx, count):
    """Positive nodes and weights of an even Gauss-Legendre rule."""
    assert count % 2 == 0
    x, w = [], []
    for k in range(1, count//2+1):
        root = ctx.cos(ctx.pi*(ctx.mpf(k)-ctx.mpf(1)/4)/(ctx.mpf(count)+ctx.mpf(1)/2))
        for step in range(30):
            p, previous = ctx.one, ctx.zero
            for j in range(1, count+1):
                p, previous = ((2*j-1)*root*p-(j-1)*previous)/j, p
            derivative = count*(root*p-previous)/(root*root-1)
            delta = p/derivative
            root -= delta
            if abs(delta) < 100*ctx.eps:
                break
        else:
            raise AssertionError('Gauss root did not converge')
        # Re-evaluate derivative at the converged root.
        p, previous = ctx.one, ctx.zero
        for j in range(1, count+1):
            p, previous = ((2*j-1)*root*p-(j-1)*previous)/j, p
        derivative = count*(root*p-previous)/(root*root-1)
        x.append(root)
        w.append(2/((1-root*root)*derivative**2))
    return x, w


def matrices(ctx, S, x, w, constant=False):
    c = ctx.pi*ctx.mpf(S)/2
    def kernel(z):
        if constant:
            return ctx.one
        return ctx.sin(c*z)/(ctx.pi*z) if z else c/ctx.pi
    even, odd = ctx.matrix(len(x)), ctx.matrix(len(x))
    for i in range(len(x)):
        for j in range(i+1):
            fac = ctx.sqrt(w[i]*w[j])
            direct, reflected = kernel(x[i]-x[j]), kernel(x[i]+x[j])
            even[i,j] = even[j,i] = fac*(direct+reflected)
            odd[i,j] = odd[j,i] = fac*(direct-reflected)
    return even, odd


def controls():
    ctx = mp.mp.clone(); ctx.dps = 100
    x,w = gauss_nodes(ctx,8)
    for degree in range(16):
        moment = 2*ctx.fsum(wi*xi**degree for xi,wi in zip(x,w)) if degree%2==0 else ctx.zero
        expected = ctx.mpf(2)/(degree+1) if degree%2==0 else ctx.zero
        assert abs(moment-expected) < ctx.mpf('1e-95')
    even,odd = matrices(ctx,1,x,w,constant=True)
    vals=sorted(list(ctx.eigsy(even,eigvals_only=True))+list(ctx.eigsy(odd,eigvals_only=True)),reverse=True)
    assert abs(vals[0]-2)<ctx.mpf('1e-95')
    assert max(abs(v) for v in vals[1:])<ctx.mpf('1e-95')
    # A nonzero rank-one control detects accidental zeroing/coercion.
    assert agreement_digits('1.0','1.01') < 3
    return {'gauss_moments_degrees':list(range(16)), 'constant_kernel_spectrum':'2,0,0,0,0,0,0,0', 'comparison_rejects_one_percent_error':True}


def nystrom(S, dps, nodes):
    ctx = mp.mp.clone(); ctx.dps=dps
    x,w=gauss_nodes(ctx,nodes)
    even,odd=matrices(ctx,S,x,w)
    vals=sorted(list(ctx.eigsy(even,eigvals_only=True))+list(ctx.eigsy(odd,eigvals_only=True)),reverse=True)
    trace=ctx.fsum(vals)
    square_trace=ctx.fsum(v*v for v in vals)
    c=ctx.pi*ctx.mpf(S)/2
    hs=(4*c*ctx.si(4*c)-(1-ctx.cos(4*c))-ctx.euler-ctx.log(4*c)+ctx.ci(4*c))/ctx.pi**2
    assert abs(trace-S)<ctx.mpf('1e-'+str(dps-15))
    assert abs(square_trace-hs)<ctx.mpf('1e-120')
    return [ctx.nstr(v,dps-10) for v in vals], {
        'trace_error':ctx.nstr(abs(trace-S),8),
        'hilbert_schmidt_error':ctx.nstr(abs(square_trace-hs),8),
        'first_12_trace_deficit':ctx.nstr(S-ctx.fsum(vals[:12]),30),
    }


def flatten_numbers(doc):
    # Read the actual API's nested parameter tree, preserving raw decimals.
    result={}
    for S, modes in doc['Numbers'].items():
        for n, row in modes.items():
            result[(int(S),int(n))] = row['number'] if isinstance(row,dict) else row
    return result


def check_stored(evidence):
    """Check database decimals against the independent pre-fill computation."""
    values = flatten_numbers(numberdb.table('T519'))
    expected = {(S,n) for S in SHANNON_NUMBERS for n in range(MODES)}
    assert set(values) == expected
    ctx = mp.mp.clone(); ctx.dps = 230
    by_key = {(row['S'],row['n']): row for row in evidence['rows']}
    rows = []
    for S,n in sorted(expected):
        text = str(values[(S,n)])
        value = ctx.mpf(text)
        mantissa,_,exponent = text.lower().partition('e')
        decimals = len(mantissa.partition('.')[2])
        ulp = ctx.power(10, int(exponent or 0)-decimals)
        reference = ctx.mpf(by_key[(S,n)]['independent_value'])
        error = abs(value-reference)/ulp
        assert 0 < value < 1 and error <= 1, (S,n,'stored interval')
        rows.append({'S':S,'n':n,'error_ulps':ctx.nstr(error,20)})
    bandwidth_checks = {}
    for S in SHANNON_NUMBERS:
        actual = [ctx.mpf(str(values[(S,n)])) for n in range(MODES)]
        reference = [ctx.mpf(by_key[(S,n)]['independent_value']) for n in range(MODES)]
        assert all(a>b for a,b in zip(actual,actual[1:]))
        trace_deficit = S-ctx.fsum(actual)
        assert trace_deficit > 0
        assert abs(ctx.fsum(actual)-ctx.fsum(reference)) < ctx.mpf('1e-98')
        assert abs(ctx.fsum(v*v for v in actual)-ctx.fsum(v*v for v in reference)) < ctx.mpf('1e-98')
        bandwidth_checks[str(S)] = {'ordered':True,'strictly_between_zero_and_one':True,'first_12_trace_deficit':ctx.nstr(trace_deficit,35)}
    return {'table':'T519','checked':len(rows),'independent_method':'120-node Nystrom at 170 decimal digits; pre-fill evidence.json','maximum_error_ulps':max(float(row['error_ulps']) for row in rows),'rows':rows,'bandwidth_checks':bandwidth_checks}


def main():
    started=time.monotonic()
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN')=='1':
        os.environ['NUMBERDB_API_KEY']=sys.stdin.read().strip()
    from sage.env import SAGE_VERSION
    evidence={'table':'T519','versions':{'python':platform.python_version(),'sage':SAGE_VERSION,'mpmath':mp.__version__,'numberdb':numberdb.__version__},'controls':controls(),'rows':[],'bandwidth_checks':{}}
    print('Controls passed.',file=sys.stderr,flush=True)
    stored_mode='--stored' in sys.argv or os.environ.get('CHECK_STORED')=='1'
    stored_values=flatten_numbers(numberdb.table('T519')) if stored_mode else None
    if stored_mode:
        assert set(stored_values)=={(S,n) for S in SHANNON_NUMBERS for n in range(MODES)}
    for S in SHANNON_NUMBERS:
        primary=[spectrum(S,dps,terms) for dps,terms in SETTINGS]
        print('Legendre complete: S='+str(S),file=sys.stderr,flush=True)
        independent,identities=nystrom(S,170,120)
        evidence['bandwidth_checks'][str(S)]=identities
        for n in range(MODES):
            agreements=[agreement_digits(a[n],b[n]) for a,b in zip(primary,primary[1:])]
            independent_digits=agreement_digits(primary[-1][n],independent[n])
            assert min(agreements)>=110,(S,n,'primary')
            assert independent_digits>=110,(S,n,'independent',str(independent_digits))
            row={'S':S,'n':n,'number':rounded(primary[-1][n],100),'precision_and_truncation_agreement':[str(x) for x in agreements],'nystrom_agreement_digits':str(independent_digits),'independent_value':independent[n]}
            if stored_mode:
                # Verify the uncertainty represented by the stored last place.
                ctx=mp.mp.clone();ctx.dps=230
                written=str(stored_values[(S,n)])
                number=ctx.mpf(written)
                mantissa,_,exponent=written.lower().partition('e')
                decimals=len(mantissa.partition('.')[2])
                ulp=ctx.power(10,int(exponent or 0)-decimals)
                ratio=abs(number-ctx.mpf(independent[n]))/ulp
                assert ratio<=1,(S,n,'outside stored interval')
                row['stored_error_ulps']=str(ratio)
                assert 0<number<1
            evidence['rows'].append(row)
        if stored_mode:
            ctx=mp.mp.clone();ctx.dps=230
            vals=[ctx.mpf(str(stored_values[(S,n)])) for n in range(MODES)]
            assert all(a>b for a,b in zip(vals,vals[1:]))
            assert ctx.fsum(vals)<S
        print('Nystrom and identities passed: S='+str(S),file=sys.stderr,flush=True)
    evidence['stored_checked']=stored_mode
    evidence['elapsed_seconds']=time.monotonic()-started
    print(json.dumps(evidence,indent=2))


if __name__=='__main__':
    main()