import sys, os; sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import sympy as sp
from grass import *

B, F, tB, tF = sp.symbols('B F theta_B theta_F', real=True)
P = sp.sqrt(B**2 - 1)
R = sp.sqrt(1 - F**2)
I = sp.I
par = [0, 1, 0, 1]   # (RB, RF, AB, AF)
X1, X2, X3, X4 = [G.gen(i) for i in (1, 2, 3, 4)]

def Q0mat():
    M = mzero(4)
    M[0][0] = G.c(B); M[1][1] = G.c(F); M[2][2] = G.c(-B); M[3][3] = G.c(-F)
    M[0][2] = G.c(P * sp.exp(I * tB)); M[1][3] = G.c(I * R * sp.exp(I * tF))
    M[2][0] = G.c(-P * sp.exp(-I * tB)); M[3][1] = G.c(-I * R * sp.exp(-I * tF))
    return M

def kblock(a, b):
    n = a * b
    return [[G.c(1) + n * sp.Rational(1, 2), a], [b, G.c(1) - n * sp.Rational(1, 2)]]

def Kmat(a1, b1, a2, b2):
    M = mzero(4)
    kR = kblock(a1, b1); kA = kblock(a2, b2)
    for i in range(2):
        for j in range(2):
            M[i][j] = kR[i][j]; M[2 + i][2 + j] = kA[i][j]
    return M

K = Kmat(X1, X2, X3, X4)
Kinv = Kmat(-X1, -X2, -X3, -X4)   # exp(-N)
Q0 = Q0mat()
Q = mmul(Kinv, mmul(Q0, K))

def simpG(g):
    return g.applyc(lambda v: sp.simplify(sp.expand(v)))

def str_(M):
    s = G()
    for i in range(4):
        s = s + (M[i][i] if par[i] == 0 else -M[i][i])
    return s

if __name__ == "__main__":
    KK = mmul(K, Kinv)
    print("K Kinv = I:", all(simpG(KK[i][j] - (G.c(1) if i == j else G())).d == {} for i in range(4) for j in range(4)))
    Q2 = mmul(Q, Q)
    print("Q^2 = I:", all(simpG(Q2[i][j] - (G.c(1) if i == j else G())).d == {} for i in range(4) for j in range(4)))
    print("str Q =", simpG(str_(Q)))
    Lam = mzero(4)
    for i, s in enumerate([1, 1, -1, -1]):
        Lam[i][i] = G.c(s)
    print("d_1 = str(Lambda Q)/2 =", simpG(str_(mmul(Lam, Q)) * sp.Rational(1, 2)))
    src = (Q[0][0] + Q[0][2] - Q[2][0] - Q[2][2]) * sp.Rational(1, 2)
    src = simpG(src)
    for k, v in sorted(src.d.items()):
        print("source coeff", k, ":", sp.simplify(v))