"""Independent checks for Zonal polynomials -- numberdb.org/T576.

Run: sage -python checks.py
Run against a saved API document: sage -python checks.py --stored stored.json
Requires generate.py and NumberDB's check.py alongside this file.

The coefficient comparison uses Sage's Jack J basis and the independent
hook-product conversion from Koev--Edelman (3.1)-(3.2). Checks outside that
construction use DLMF 35.4.2 at the identity and the two-dimensional Haar
integral in DLMF 35.4.3, evaluated by exact even trigonometric moments.
Jiu--Koutschan Example 3.5 is an independently transcribed degree-four table.
No mathematical routine from generate.py is used to compute expected values.
"""
import numberdb.sage as numberdb
from sage.rings.rational_field import QQ
from sage.rings.polynomial.polynomial_ring_constructor import PolynomialRing
from sage.symbolic.ring import SR
from sage.combinat.sf.sf import SymmetricFunctions
from sage.combinat.partition import Partitions
from sage.env import SAGE_VERSION
from math import factorial, comb
from itertools import combinations
from collections import Counter
import json
import sys
from pathlib import Path
from generate import ZonalPolynomials, DEGREE_LIMITS
from check import exactness, measure, agrees_with, names_its_rings

J = SymmetricFunctions(QQ).jack(t=QQ(2)).J()


def new_ring(n):
    return PolynomialRing(QQ,n,names=['x%d' % i for i in range(1,n+1)])


def product(seq):
    answer=QQ(1)
    for x in seq: answer *= x
    return answer


def hook_product(lam):
    return product((2*(row-j)+sum(a>=j for a in lam)-i+1)*
                   (2*(row-j+1)+sum(a>=j for a in lam)-i)
                   for i,row in enumerate(lam,1) for j in range(1,row+1))


def jack_expected(key):
    n,lam=key;d=sum(lam);R=new_ring(n)
    f=R(str(J[list(lam)].expand(n,alphabet=list(R.variable_names()))))
    return QQ(2**d*factorial(d))/hook_product(lam)*f


def at_identity(n,lam):
    """DLMF 35.4.2, using row differences, not hook products."""
    d=sum(lam);ell=len(lam)
    shifted=product(QQ(n-i+1)/QQ(2)+j for i,row in enumerate(lam,1) for j in range(row))
    numerator=product(2*lam[i]-2*lam[j]-i+j for i in range(ell) for j in range(i+1,ell))
    denominator=product(factorial(2*row+ell-i) for i,row in enumerate(lam,1))
    return QQ(factorial(d)*4**d)*shifted*numerator/denominator


def circle_moment(k,l):
    return QQ(factorial(2*k)*factorial(2*l))/QQ(4**(k+l)*factorial(k)*factorial(l)*factorial(k+l))


def haar_two_variables(lam):
    a=lam[0];b=lam[1] if len(lam)==2 else 0;r=a-b
    R=new_ring(2);x,y=R.gens()
    average=sum((QQ(comb(r,k))*circle_moment(k,r-k)*x**k*y**(r-k) for k in range(r+1)),R.zero())
    return at_identity(2,lam)*(x*y)**b*average


# Jiu--Koutschan, Example 3.5, coefficient rows and columns in this order.
D4_SHAPES=((4,),(3,1),(2,2),(2,1,1),(1,1,1,1))
D4_COEFFICIENTS=(
    ('1','4/7','18/35','12/35','8/35'),
    ('0','24/7','16/7','88/21','32/7'),
    ('0','0','16/5','32/15','16/5'),
    ('0','0','0','16/3','64/5'),
    ('0','0','0','0','16/5'))


def controls():
    assert circle_moment(0,0)==1
    assert circle_moment(1,0)==QQ(1)/2
    assert circle_moment(1,1)==QQ(1)/8
    assert at_identity(2,(2,))==QQ(8)/3
    assert at_identity(2,(1,1))==QQ(4)/3
    R=new_ring(2);x,y=R.gens()
    assert jack_expected((2,(2,)))==x*x+QQ(2)/3*x*y+y*y
    assert haar_two_variables((2,))==x*x+QQ(2)/3*x*y+y*y
    assert haar_two_variables((1,1))==QQ(4)/3*x*y
    return 8


def read_stored(path):
    """Convert nested or flat NumberDB entries without relying on generator code."""
    doc=json.loads(Path(path).read_text());block=doc['Numbers'];values={}
    if isinstance(block,list):
        rows=block
    else:
        rows=[]
        for n,by_shape in block.items():
            for lam,entry in by_shape.items():
                rows.append({'params':{'n':str(n),'lambda':str(lam)},**(entry if isinstance(entry,dict) else {'number':entry})})
    for row in rows:
        p=row['params'];n=int(p['n']);lam=tuple(map(int,str(p['lambda']).split(',')))
        assert (n,lam) not in values
        values[n,lam]=new_ring(n)(row['number'])
    return values


def run(values):
    counts={'controls':controls()}
    # Enumeration uses Sage's independent partition iterator.
    expected=[(n,tuple(map(int,lam))) for n,dmax in DEGREE_LIMITS for d in range(2,dmax+1)
              for lam in Partitions(d,max_length=n)]
    assert list(values)==expected,'stored/default order or selection differs'
    counts['enumerated']=len(expected)
    assert not exactness(values)
    assert not names_its_rings(str(Path(__file__).with_name('generate.py')))
    differences=agrees_with(values,jack_expected)
    assert not differences,differences
    counts['sage_jack_full_coefficients']=len(values)
    for (n,lam),f in values.items():
        d=sum(lam);R=f.parent();x=R.gens();terms={tuple(e):c for e,c in f.dict().items()}
        assert f(*([1]*n))==at_identity(n,lam)
        assert all(sum(e)==d for e in terms)
        # Adjacent swaps generate every permutation of the alphabet.
        for i in range(n-1):
            swap=list(x);swap[i],swap[i+1]=swap[i+1],swap[i]
            assert f(*swap)==f
        pairs=list(combinations(range(n),2))
        delta=product(x[i]-x[j] for i,j in pairs)
        eigenvalue=sum(a*(a+n-i-1) for i,a in enumerate(lam,1))
        residual=delta*(sum((v*v*f.derivative(v,2) for v in x),R.zero())-eigenvalue*f)
        for i,j in pairs:
            others=product(x[a]-x[b] for a,b in pairs if (a,b)!=(i,j))
            residual += others*(x[i]**2*f.derivative(x[i])-x[j]**2*f.derivative(x[j]))
        assert residual==0
        if n==2: assert f==haar_two_variables(lam)
        # Published coefficient table, distinct from both executable algorithms.
        if d==4:
            row=D4_COEFFICIENTS[D4_SHAPES.index(lam)]
            for mu,c in zip(D4_SHAPES,row):
                if len(mu)<=n: assert terms.get(mu+(0,)*(n-len(mu)),QQ(0))==QQ(c)
        if len(lam)==d:
            elementary=sum((product(x[i] for i in indices) for indices in combinations(range(n),d)),R.zero())
            assert f==QQ(2**d)/QQ(d+1)*elementary
        if d==2:
            p1=sum(x);p2=sum(a*a for a in x)
            wanted=(p1*p1+2*p2)/QQ(3) if lam==(2,) else QQ(2)*(p1*p1-p2)/QQ(3)
            assert f==wanted
        if n>2:
            specialized=f(*(list(x[:-1])+[R.zero()]))
            if len(lam)==n: assert specialized==0
            else:
                lower=values.get((n-1,lam),jack_expected((n-1,lam)))
                assert specialized==R(str(lower))
    counts['identity_value']=len(values)
    counts['cleared_denominator_eigenidentity']=len(values)
    counts['symmetry_and_homogeneity']=len(values)
    counts['haar_integral_n2']=sum(n==2 for n,lam in values)
    counts['published_degree4_rows']=sum(sum(lam)==4 for n,lam in values)
    counts['degree2_formulas']=sum(sum(lam)==2 for n,lam in values)
    counts['column_formula']=sum(len(lam)==sum(lam) for n,lam in values)
    counts['alphabet_stability']=sum(n>2 for n,lam in values)
    for n,dmax in DEGREE_LIMITS:
        R=new_ring(n)
        for d in range(2,dmax+1):
            assert sum((f for (nn,lam),f in values.items() if nn==n and sum(lam)==d),R.zero())==sum(R.gens())**d
    counts['trace_power_slices']=sum(dmax-1 for n,dmax in DEGREE_LIMITS)
    # Boundary identities are prose claims; check them for all selected degrees.
    for n in range(1,6):
        R=new_ring(n);x=R.gens()
        assert jack_expected((n,()))==1
        assert jack_expected((n,(1,)))==sum(x)
    for d in range(1,7):
        R=new_ring(1);assert jack_expected((1,(d,)))==R.gen()**d
    counts['boundary_checks']=16
    print(json.dumps({'passed':counts,'measure':measure({str(k):v for k,v in values.items()}),
                      'versions':{'sage':SAGE_VERSION,'numberdb':getattr(numberdb,'__version__','unknown')}},indent=2))
    return counts


if __name__=='__main__':
    if '--stored' in sys.argv:
        values=read_stored(sys.argv[sys.argv.index('--stored')+1])
    else:
        g=ZonalPolynomials();values={(int(p['n']),tuple(map(int,p['lambda'].split(',')))):g.value(p,100) for p in g.enumerate()}
    run(values)
