"""Independent checks for Witten zeta values -- numberdb.org/T607.

    sage -python checks.py

Optional: sage -python checks.py document.json checks-evidence.json
The JSON document must be an API table document with flat or nested Numbers records.
No writing API is used. Without a document, checks the generated values.
Reads generate.py; exact published coefficients in this file are independent
of its Weyl-lattice computation. Sources and error bounds are in derivation.md.
"""
import json
import sys
from decimal import Decimal
from fractions import Fraction
from math import factorial, comb
from pathlib import Path
import numberdb.sage as numberdb
from numberdb._generate import _stored_entries
from sage.rings.rational_field import QQ
from sage.rings.real_arb import RealBallField
from sage.rings.polynomial.polynomial_ring_constructor import PolynomialRing
from generate import WittenZeta, coefficient, partial_fractions

# Euler's first four even zeta values; Zhao (2010), equation (1.1);
# Zhao (2009), MPIM 2009-42, section 4, page 8 (pi-power forms);
# Komori-Matsumoto-Tsumura IV, Example 2.2, page 5.
A1 = [Fraction(1,6),Fraction(1,90),Fraction(1,945),Fraction(1,9450)]
B2_RAW = [Fraction(6,5*factorial(9)),Fraction(2**5*479,5*factorial(17)),
          Fraction(2**7*5*43*19309,3**2*7*13*factorial(23)),
          Fraction(2**8*13*241*64009163,5*17*factorial(31))]
G2_RAW = [Fraction(23,297904566960),
          Fraction(8165653,1445838676129559305994400000),
          Fraction(55940539974690617,131888156302530666544150214880458495963616000000),
          Fraction(47346365461279256768015189,1485697621623958244738368714652675148113575302412190275200000000000)]

# A second Bernoulli algorithm, triangular rather than the generator recurrence.
def zeta_ratio(n):
    if n==0:
        return Fraction(-1,2)
    a=[Fraction(0)]*(n+1)
    for m in range(n+1):
        a[m]=Fraction(1,m+1)
        for j in range(m,0,-1):
            a[j-1]=j*(a[j-1]-a[j])
    return (-1)**(n//2+1)*a[0]*Fraction(2**(n-1),factorial(n))

def published_coefficient(key):
    algebra,m=key
    if algebra=='A1': return QQ(A1[m-1]),2*m
    if algebra=='A2':
        c=Fraction(4**(m+1),3)*sum(
            comb(4*m-i-1,2*m-1)*zeta_ratio(i)*zeta_ratio(6*m-i)
            for i in range(0,2*m+1,2))
        return QQ(c),6*m
    if algebra=='B2':return QQ(B2_RAW[m-1]*6**(2*m)),8*m
    if algebra=='G2':return QQ(G2_RAW[m-1]*120**(2*m)),12*m
    raise ValueError(algebra)

def decimal_interval(text):
    d=Decimal(text)
    unit=Decimal(1).scaleb(d.as_tuple().exponent)
    return QQ(Fraction(d))-QQ(Fraction(unit)),QQ(Fraction(d))+QQ(Fraction(unit))

def finite_positive_sums(algebra,bound=256):
    field=RealBallField(600)
    sums=[field(0) for _ in range(4)]
    for a in range(1,bound+1):
        for b in range(1,bound+1):
            if algebra=='A2': dim=QQ(a*b*(a+b))/2
            elif algebra=='B2': dim=QQ(a*b*(a+b)*(a+2*b))/6
            elif algebra=='G2': dim=QQ(a*b*(a+b)*(a+2*b)*(a+3*b)*(2*a+3*b))/120
            else:raise ValueError(algebra)
            assert dim.denominator()==1 and dim>=1
            inverse=field(dim)**(-2)
            power=inverse
            for i in range(4):
                sums[i]+=power
                power*=inverse
    # P(a,b) >= Ca*a**(N-1)*b and >= Cb*a*b**(N-1).
    N,M,Ca,Cb={'A2':(3,2,1,1),'B2':(4,6,1,2),'G2':(6,120,2,18)}[algebra]
    tails=[]
    for m in range(1,5):
        s=2*m; p=(N-1)*s
        tails.append(field(M)**s*(field(Ca)**(-s)+field(Cb)**(-s))*field(s).zeta()*field(bound)**(1-p)/field(p-1))
    return sums,tails

def run(document=None):
    expected=[(a,m) for a in ('A1','A2','B2','G2') for m in range(1,5)]
    g=WittenZeta();params=list(g.enumerate())
    assert [(p['algebra'],int(p['m'])) for p in params]==expected
    field=RealBallField(600)
    # Known-answer controls, including nonzero values, before numerical tests.
    assert zeta_ratio(2)==Fraction(1,6)
    assert (field(2).zeta()-field.pi()**2/6).contains_zero()
    assert field(2).zeta().is_finite() and field(2).zeta().rad_as_ball()<field(10)**-170
    control_sum=sum(field(k)**-2 for k in range(1,257))
    assert control_sum < field(2).zeta() < control_sum+field(1)/256
    assert not control_sum < field(2).zeta()+1 < control_sum+field(1)/256
    qring=PolynomialRing(QQ,'x');x=qring.gen();fraction_field=qring.fraction_field()
    exact_poles={'A2':(x,x+1),'B2':(x,x+1,x+2),'G2':(x,x+1,x+2,x+3,2*x+3)}
    for algebra in ('A2','B2','G2'):
        for m in range(1,5):
            denominator=qring(1)
            for term in exact_poles[algebra]:denominator*=term**(2*m)
            actual=sum(fraction_field(c)/(x-r)**k for r,k,c in partial_fractions(algebra,2*m))
            assert actual==1/fraction_field(denominator)
    finite={a:finite_positive_sums(a) for a in ('A2','B2','G2')}
    stored={}
    if document:
        assert document['Title']=='Witten zeta values of simple Lie algebras'
        written=_stored_entries(document)
        identities=[(key.split(',')[0],int(key.split(',')[1])) for key in written]
        assert identities==expected
        stored=dict(zip(identities,written.values()))
    rows=[]
    for p,key in zip(params,expected):
        algebra,m=key
        c,w=coefficient(algebra,m);pc,pw=published_coefficient(key)
        assert c==pc and w==pw
        got=g.value(p,100)
        independently=field(pc)*field.pi()**pw
        if algebra=='A1':independently=field(2*m).zeta()
        if algebra=='A2':
            independently=field(4)**(m+1)/3*sum(field(comb(4*m-i-1,2*m-1))*field(i).zeta()*field(6*m-i).zeta() for i in range(0,2*m+1,2))
        assert got.is_finite() and independently.is_finite()
        assert got.rad_as_ball()<field(10)**-120 and independently.rad_as_ball()<field(10)**-170
        assert got.overlaps(independently)
        assert not (got+field(10)**-90).overlaps(independently)
        row={'algebra':algebra,'m':m,'coefficient':str(c),'pi_exponent':w,
             'accuracy_bits':int(got.accuracy()),'published_exact_agreement':True,
             'independent_ball_agreement':True}
        if algebra in finite:
            sums,tails=finite[algebra];low=sums[m-1];tail=tails[m-1]
            assert low < got < low+tail, (key, str(low), str(got), str(tail), str(got-low))
            assert tail<field(10)**-5
            assert not low < got+field(10)**-4 < low+tail
            row.update(finite_sum_bound=256,finite_sum=str(low),tail_bound=str(tail),tail_contains_value=True)
        if key in stored:
            lo,hi=decimal_interval(stored[key])
            assert lo<independently and independently<hi
            assert not lo<independently+10*(hi-lo)<hi
            row['stored_interval_contains_independent_ball']=True
        rows.append(row)
    for algebra in ('A1','A2','B2','G2'):
        vals=[g.value({'algebra':algebra,'m':str(m)},100) for m in range(1,5)]
        assert all(vals[i]>vals[i+1]>1 for i in range(3))
        if stored:
            intervals=[decimal_interval(stored[(algebra,m)]) for m in range(1,5)]
            assert all(intervals[i][0]>intervals[i+1][1]>1 for i in range(3))
    # G2 sum uses the same dimensions as the 7- and 14-dimensional fundamental representations.
    assert 2*1*3*4*5*7//120==7 and 1*2*3*5*7*8//120==14
    return {'table':'T607','skill_manifest':'agents/runs/20261029T103821Z-build-skills/manifest.json',
            'skill_sha256':'28ecbeeccb0c5e5d6eb541b5a5b0f88aec47b154374d909085116d34f5629199',
            'entry_count':16,'controls_passed':True,'partial_fraction_identities_checked':12,
            'minimum_accuracy_bits':min(r['accuracy_bits'] for r in rows),
            'stored_values_checked':bool(document),'rows':rows}

if __name__=='__main__':
    document=json.loads(Path(sys.argv[1]).read_text()) if len(sys.argv)>1 else None
    result=run(document)
    output=json.dumps(result,indent=2)+'\n'
    if len(sys.argv)>2:Path(sys.argv[2]).write_text(output)
    print(output)