import sys, os; sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import sympy as sp, pickle, time
from grass import *
from chart import B, F, tB, tF, P, R, Q, par

even_coords = [B, F, tB, tF]
odd_gens = [1, 2, 3, 4]
ETA = 5
ent_even = [(0, 0), (1, 1), (0, 2), (1, 3)]
ent_odd = [(0, 1), (1, 0), (2, 3), (3, 2)]
ents = ent_even + ent_odd

def dI(f, I):
    return f.diff(even_coords[I]) if I < 4 else f.dleft(odd_gens[I - 4])

# Jacobian J[I][e] = d_I Q_e
J = [[dI(Q[e[0]][e[1]], I) for e in ents] for I in range(8)]
J0 = sp.Matrix(8, 8, lambda I, e: J[I][e].body())
J0inv = sp.simplify(J0.inv())
Jn = [[J[I][e] - G.c(J[I][e].body()) for e in range(8)] for I in range(8)]
J0invG = [[G.c(J0inv[r, c]) for c in range(8)] for r in range(8)]
M = mscale(-1, mmul(J0invG, Jn))
S = meye(8); Mk = meye(8)
for k in range(1, 7):
    Mk = mmul(Mk, M)
    S = madd(S, Mk)
Jinv = mmul(S, J0invG)     # rows e, cols I

def vfield(A):
    """components v^I of the vector field with dQ = [A, Q] for an even supermatrix A"""
    dQ = msub(mmul(A, Q), mmul(Q, A))
    row = [dQ[e[0]][e[1]] for e in ents]
    v = []
    for I in range(8):
        s = G()
        for e in range(8):
            s = s + row[e] * Jinv[e][I]
        v.append(s)
    # consistency on all 16 entries
    return v, dQ

def apply_v(v, f):
    out = G()
    for I in range(8):
        if v[I].d:
            out = out + v[I] * dI(f, I)
    return out

basis = []
t0 = time.time()
for i in range(4):
    for j in range(4):
        p = (par[i] + par[j]) % 2
        A = mzero(4)
        A[i][j] = G.c(1) if p == 0 else G.gen(ETA)
        v, dQ = vfield(A)
        basis.append(((i, j), p, v))
print("vector fields built in", time.time() - t0)

def X(ij, f):
    for (kl, p, v) in basis:
        if kl == ij:
            g = apply_v(v, f)
            return g if p == 0 else g.dleft(ETA)

def casimir(f):
    """Return the auxiliary quadratic operator -2*C_1; callers normalize it."""
    out = G()
    for i in range(4):
        for j in range(4):
            sgn = (-1) ** par[j]
            if (par[i] + par[j]) % 2 == 1:
                # e:sigma-casimir after renaming i <-> j: sum (-1)^{|i|} E_ij E_ji;
                # for odd E_ij, (-1)^{|i|} = -(-1)^{|j|}
                sgn = -sgn
            out = out + X((i, j), X((j, i), f)) * sgn
    return out

if __name__ == "__main__":
    # consistency check of one odd vector field on all entries
    (ij, p, v) = basis[1]
    A = mzero(4); A[0][1] = G.gen(ETA)
    dQ = msub(mmul(A, Q), mmul(Q, A))
    import random
    pt = {B: sp.Rational(23, 10), F: sp.Rational(3, 10), tB: sp.Rational(7, 10), tF: sp.Rational(-4, 10)}
    bad = 0
    for k in range(4):
        for l in range(4):
            lhs = apply_v(v, Q[k][l]) - dQ[k][l]
            for mono, c in lhs.d.items():
                if abs(complex(sp.N(c.subs(pt), 30))) > 1e-20:
                    bad += 1
    print("tangency check (odd field E_01) mismatches:", bad)
    with open('basis.pkl', 'wb') as fh:
        pickle.dump([(ij, p, v) for (ij, p, v) in basis], fh)