"""Independent geometric checks of T498. Run: sage -python checks.py.

No numerical formula from generate.py is used by independent_form().
Root cells: integrate the quadratic form over an affine Weyl alcove.
A* cells: integrate pyramids over product-permutohedron facets, without
Conway-Sloane's Abel-identity simplification used in the generator.
D* cells: integrate one coordinate squared against the Irwin-Hall CDF,
using exact polynomial antiderivatives on each clipped interval.
Each check is exact before taking the positive nth root in Arb.
"""
import os
import sys
import json
from functools import lru_cache
from math import factorial, comb
import numberdb.sage as numberdb
from sage.rings.rational_field import QQ
from sage.rings.real_arb import RealBallField
from generate import LatticeSecondMoments, exact_form, evaluate, permutohedron_j
from check import exactness, measure, agrees_with, names_its_rings


def dot(u, v):
    return sum(a*b for a,b in zip(u,v))


def simplex_mean(vertices):
    # Includes the zero vertex. Barycentric expectations:
    # E[t_i^2]=2/((n+1)(n+2)), E[t_i*t_j]=1/((n+1)(n+2)).
    n = len(vertices)-1
    total = [sum(v[i] for v in vertices) for i in range(len(vertices[0]))]
    return (sum(dot(v,v) for v in vertices)+dot(total,total))/((n+1)*(n+2))


def inverse(a):
    n=len(a)
    rows=[[QQ(x) for x in row]+[QQ(i==j) for j in range(n)] for i,row in enumerate(a)]
    det=QQ(1)
    for i in range(n):
        k=next(j for j in range(i,n) if rows[j][i])
        if k!=i:
            rows[i],rows[k]=rows[k],rows[i]; det=-det
        p=rows[i][i]; det*=p
        rows[i]=[x/p for x in rows[i]]
        for j in range(n):
            if j!=i:
                p=rows[j][i]
                rows[j]=[x-p*y for x,y in zip(rows[j],rows[i])]
    return [r[n:] for r in rows],det


def exceptional_mean(n):
    # E_n diagram: chain of n-1 vertices, with the final vertex attached
    # at chain index 2. Arms have lengths 2, n-4, 1.
    c=[[2*int(i==j) for j in range(n)] for i in range(n)]
    for i,j in [(i,i+1) for i in range(n-2)]+[(2,n-1)]:
        c[i][j]=c[j][i]=-1
    roots={tuple(int(j==i) for j in range(n)) for i in range(n)}
    todo=list(roots)
    while todo:
        r=todo.pop()
        for i in range(n):
            s=list(r); s[i]-=sum(c[i][j]*r[j] for j in range(n)); s=tuple(s)
            if s not in roots: roots.add(s); todo.append(s)
    assert len(roots)=={6:72,7:126,8:240}[n]
    marks=max(roots,key=sum)
    inv,det=inverse(c)
    assert det=={6:3,7:2,8:1}[n]
    # Vertices fundamental_weight_i / highest_root_coefficient_i;
    # pairwise Euclidean products are inv[i][j]/(marks[i]*marks[j]).
    mean=(sum(inv[i][i]/marks[i]**2 for i in range(n))+
          sum(inv[i][j]/(marks[i]*marks[j]) for i in range(n) for j in range(n))) / ((n+1)*(n+2))
    return mean,det


@lru_cache(None)
def permutohedron_mean(d):
    if d==0: return QQ(0)
    total=QQ(0)
    for r in range(d):
        s=d-1-r
        height2=QQ(((r+1)*(s+1)*(d+1),4))
        factor=QQ((comb(d+1,r+1)*(r+1)**r*(s+1)**s,
                  2*(d+2)*(d+1)**(d-1)))
        total+=factor*(height2+permutohedron_mean(r)+permutohedron_mean(s))
    return total


def clipped_cube_integral(n):
    # Integrate y^2 times vol{z in [0,1]^(n-1): sum z <= n/2-y}.
    # Multiply by n for coordinate symmetry, then divide by 4 for x=y/2.
    total=QQ(0)
    for j in range(n):
        a=QQ((n,2))-j
        if a<=0: continue
        b=min(QQ(1),a)
        integ=sum((-1)**k * comb(n-1,k)*a**(n-1-k)*b**(k+3)/QQ(k+3)
                  for k in range(n))
        total+=(-1)**j*QQ((comb(n-1,j),factorial(n-1)))*integ
    return QQ((n,4))*total


@lru_cache(None)
def independent_form(f,n):
    if f=='Z':
        # Product of intervals: integral norm^2 = n * integral_-1/2^1/2 x^2 dx.
        return QQ((1,12)),QQ(1)
    if f=='A':
        vertices=[[QQ(0)]*(n+1)]
        for i in range(1,n+1):
            j=n+1-i
            vertices.append([QQ((-j,n+1))]*i+[QQ((i,n+1))]*j)
        return simplex_mean(vertices)/n,QQ((1,n+1))
    if f=='D':
        vertices=[[QQ(0)]*n,[QQ(1)]+[QQ(0)]*(n-1)]
        for k in range(2,n-1): vertices.append([QQ((1,2))]*k+[QQ(0)]*(n-k))
        vertices.extend([[QQ((1,2))]*(n-1)+[QQ((-1,2))],[QQ((1,2))]*n])
        assert len(vertices)==n+1
        return simplex_mean(vertices)/n,QQ((1,4))
    if f=='E':
        mean,det=exceptional_mean(n)
        return mean/n,1/det
    if f=='A*':
        return permutohedron_mean(n)/(n*(n+1)**2),QQ(n+1)
    if f=='D*':
        return 2*clipped_cube_integral(n)/n,QQ(4)
    raise ValueError(f)


def main():
    # Known controls precede comparisons, and a perturbed control must fail.
    assert simplex_mean([[QQ((-1,2))],[QQ((1,2))]])==QQ((1,12))
    assert simplex_mean([[QQ(0),QQ(0)],[QQ(1),QQ(0)],[QQ(0),QQ(1)]])==QQ((1,3))
    assert permutohedron_mean(1)==QQ((1,6))
    assert clipped_cube_integral(1)==QQ((1,96))
    assert clipped_cube_integral(2)==QQ((1,24))
    assert independent_form('E',8)[0]==QQ((929,12960))
    assert independent_form('A',2)[0]!=QQ((5,36))+QQ((1,10**50))
    print('Controls: interval, triangle, segment permutohedron, D1* and D2* cells, E8; perturbed value rejected.')
    g=LatticeSecondMoments(); values={}
    worst=0.0
    exact_count=0
    for p in g.enumerate():
        f,n=p['family'],int(p['n'])
        a,b=exact_form(f,n); c,d=independent_form(f,n)
        assert a**n*b==c**n*d, (p,a,b,c,d)
        x=evaluate(f,n,100); y=evaluate(f,n,150)
        if x in QQ and type(x).__name__=='Rational':
            assert x==y; exact_count+=1
        else:
            assert x.is_finite() and y.is_finite()
            assert x.overlaps(y)
            rel=float(x.rad()/abs(x.mid()))
            assert rel<1e-110
            worst=max(worst,rel)
        values[(f,n)]=g.value(p,100)
    assert not exactness(values)
    assert not names_its_rings(os.path.join(os.path.dirname(__file__),'generate.py'))
    # Shared agrees_with compares exact quantities, without hiding an error in a wide ball.
    powers={(p['family'],int(p['n'])): exact_form(p['family'],int(p['n']))[0]**int(p['n'])*exact_form(p['family'],int(p['n']))[1] for p in g.enumerate()}
    failures=agrees_with(powers,lambda key: independent_form(*key)[0]**key[1]*independent_form(*key)[1])
    assert not failures,failures
    print('Independent exact geometric comparisons:',len(values))
    print('Exact rational entries:',exact_count)
    print('100/150-digit overlap checks: all passed; worst relative radius:',worst)
    print('MEASURE',json.dumps(measure(values),default=str))
    # Published source values not used to construct independent_form.
    B=RealBallField(500)
    known=[('A',2,B(5)/(36*B(3).sqrt())),('A*',3,B(19)/(192*B(2)**QQ((1,3)))),
           ('D',4,B(13)/(120*B(2).sqrt())),('D*',5,B(2641)/(23040*B(2)**QQ((3,5)))),
           ('E',8,B(QQ((929,12960))))]
    for f,n,y in known:
        assert B(evaluate(f,n,120)).overlaps(y),(f,n)
    for n,expected in enumerate([QQ(0),QQ((1,12)),QQ((5,18)),QQ((19,32)),QQ((389,375)),QQ((1045,648)),QQ((78077,33614))],1):
        assert permutohedron_j(n)==expected,(n,permutohedron_j(n),expected)
    print('Source controls: five named lattice values and seven published permutohedron moments passed.')
    return values


if __name__=='__main__':
    main()