"""Exact symbolic check of l:res-polar (Section 10).

The difference D = (half the auxiliary Casimir of casimir.py) - (right side of e:res-polar)
is a second-order differential operator in the even coordinates (B, F, theta_B, theta_F)
with Grassmann-valued coefficients.  Such an operator vanishes identically if it vanishes
on f = xi^I h for all 16 monomials xi^I and all h in {1, x_i, x_i x_j}, where x runs over
the four even coordinates (D(1) = c, D(x_i) = b^i + c x_i, D(x_i x_j) = 2 a^{(ij)} + ...).
Each coefficient of D(f) is a function of (B, F, theta_B, theta_F); we decide that it is
identically zero by the rational substitution
   B = (t + 1/t)/2, sqrt(B^2 - 1) = (t - 1/t)/2   (t > 1),
   F = (1 - u^2)/(1 + u^2), sqrt(1 - F^2) = 2u/(1 + u^2)   (u > 0),
   exp(i theta_B) = wB, exp(i theta_F) = wF,
after which every coefficient is a rational function of (t, u, wB, wF, theta_B, theta_F)
and sympy's cancel decides zero exactly.
"""
import sys, os; sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import sympy as sp, itertools, time
from grass import *
from casimir import casimir, B, F, tB, tF
I = sp.I
p = B**2 - 1; q = 1 - F**2; d = B - F
sq = sp.sqrt(p) * sp.sqrt(q)
ph = sp.exp(I*(tB - tF))
kap = {(1, 2): 2*(B*F - 1)/d**2, (3, 4): 2*(B*F - 1)/d**2,
       (1, 4): 2*I*sq*ph/d**2, (2, 3): -2*I*sq/ph/d**2}
cpart = {1: (G.gen(2), sp.Rational(1, 2)*I), 2: (G.gen(1), sp.Rational(1, 2)*I),
         3: (G.gen(4), -sp.Rational(1, 2)*I), 4: (G.gen(3), -sp.Rational(1, 2)*I)}
def Theta(f): return f.diff(tB) + f.diff(tF)
def D(a, f):
    xi, c = cpart[a]
    return f.dleft(a) + (xi * Theta(f)) * c
def C0(f):
    return (f.diff(B) * p).diff(B) + (f.diff(F) * q).diff(F) + f.diff(B) * (-2*p/d) + f.diff(F) * (2*q/d)
def Lhyp(f):
    out = C0(f) + f.diff(tB).diff(tB) * (1/p) + f.diff(tF).diff(tF) * (1/q)
    for (a, b), k in kap.items():
        out = out + (D(b, D(a, f)) - D(a, D(b, f))) * (k/2)
    return out

t, u = sp.symbols('t u', positive=True)
wB, wF = sp.symbols('wB wF')
Pt = (t - 1/t) / 2
Ru = 2*u / (1 + u**2)
Bt = (t + 1/t) / 2
Fu = (1 - u**2) / (1 + u**2)

def ratsub(e):
    e = sp.sympify(e)
    def fix_pow(x):
        if isinstance(x, sp.Pow):
            b, ex = x.args
            if ex.is_Rational and ex.q == 2:
                bb = sp.expand(b)
                if sp.expand(bb - (B**2 - 1)) == 0:
                    return sp.Pow(Pt, 2*ex)
                if sp.expand(bb - (1 - F**2)) == 0:
                    return sp.Pow(Ru, 2*ex)
                if sp.expand(bb - (B**2 - 1)*(1 - F**2)) == 0:
                    return sp.Pow(Pt*Ru, 2*ex)
                if sp.expand(bb + (B**2 - 1)*(1 - F**2)) == 0 or sp.expand(bb - (B**2-1)*(F**2-1)) == 0:
                    raise ValueError('unexpected sign under root: %s' % b)
        return None
    e = e.replace(lambda x: fix_pow(x) is not None, fix_pow)
    def fix_exp(x):
        a = sp.expand(x.args[0] / sp.I)
        cB = a.coeff(tB); cF = a.coeff(tF)
        rest = sp.expand(a - cB*tB - cF*tF)
        return wB**cB * wF**cF * sp.exp(sp.I*rest)
    e = e.replace(lambda x: isinstance(x, sp.exp), fix_exp)
    e = e.subs({B: Bt, F: Fu})
    return e

def Ltrue(f): return casimir(f) * sp.Rational(1, 2)

if __name__ == '__main__':
    xs = [B, F, tB, tF]
    hs = [sp.Integer(1)] + xs + [xs[i]*xs[j] for i in range(4) for j in range(i, 4)]
    gens = [1, 2, 3, 4]
    monos = [()] + [m for r in range(1, 5) for m in itertools.combinations(gens, r)]
    t0 = time.time(); bad = 0; n = 0
    for mono in monos:
        for h in hs:
            f = G({mono: h})
            diff = Ltrue(f) - Lhyp(f)
            for k, v in diff.d.items():
                r = sp.expand(sp.numer(sp.together(ratsub(v))))
                if r != 0:
                    bad += 1
                    print('NONZERO', mono, h, k, r)
            n += 1
        print('monomial', mono, 'done', round(time.time() - t0, 1), 's', flush=True)
    print('checked', n, 'test functions;', bad, 'nonzero coefficients')