#sample script: calculate Riemann tensor and connection
#----------by Zhiqi Huang--for the course ''General Relativity''-------------

import sympy as sym
sym.init_printing()

#dimension of the spacetime
dim = 4

##allocate space to save connections, Riemann tensor, and Ricci tensor
#\Gamma_{ijk}
gam_down = sym.MutableDenseNDimArray(range(dim**3), shape=(dim, dim, dim))
#\Gamma^i_{ jk}
gam_up = sym.MutableDenseNDimArray(range(dim**3), shape=(dim, dim, dim))
#R_{ij}
ricci_down = sym.MutableDenseNDimArray(range(dim**2), shape=(dim, dim))
#################USER SECTION ########################
###Please define the coordinates (and their names) and metric here##########
#define the coordinates u[0],u[1],...
u = sym.symarray('u',dim)
coor_names = sym.symbols('v, u, x, y')  #the names of the coordinates
#--------------define the metric -----------------
f = sym.symbols('f', cls=sym.Function)  #define an arbitray function
lam = sym.Symbol('lambda', positive = True)  #define a parameter
gdown = sym.Matrix([[0, sym.Rational(-1, 2), 0, 0],[sym.Rational(-1, 2), -f(u[1], u[2], u[3]),0, 0],[0, 0, -1, 0],[0, 0, 0, -1]])
#######################################################################################


coor_mapping = {u[i]: coor_names[i] for i in range(dim)}

def print_regular(expr):
    print(expr.subs(coor_mapping))


def print_tex(expr):
    print(sym.latex(expr.subs(coor_mapping)))
    
#contravariant metric
gup = gdown.inv() 

def connection_down(i, j, k):
    return (sym.diff(gdown[i, j], u[k]) + sym.diff(gdown[i, k], u[j]) - sym.diff(gdown[j, k], u[i]))/2


#compute connection \Gamma_{ijk}
for i in range(dim):
    for j in range(dim):
        for k in range(j+1):
            gam_down[i,j,k] = connection_down(i, j, k) 
            if(j != k):
                gam_down[i,k,j] = gam_down[i,j,k]


def connection_up(i, j, k):
    gam = 0
    for l in range(dim):
        gam += gam_down[l,j, k] * gup[l, i]
    return sym.simplify(gam)

#compute connection \Gamma^i_{ jk}
for i in range(dim):
    for j in range(dim):
        for k in range(j+1):
            gam_up[i,j,k] = connection_up(i, j, k) 
            if(j != k):
                gam_up[i,k,j] = gam_up[i,j,k]

##===================compute Ricci tensor ==============================
def Riemann_tensor_up(i, j, k, l): ## R^i_{  jkl}
    R = sym.diff(gam_up[i, j, l], u[k]) - sym.diff(gam_up[i, j, k], u[l])
    for m in range(dim):
        R += gam_up[i, m, k] * gam_up[m, j, l] - gam_up[i, m, l] * gam_up[m, j, k] 
    return sym.simplify(R)

for i in range(dim):
    for j in range(i+1):
        ricci_down[i, j] = 0
        for k in range(dim):
            ricci_down[i, j] += Riemann_tensor_up(k, i, k, j)
        ricci_down[i,j] = sym.simplify(ricci_down[i,j])
        if(i != j):
            ricci_down[j, i] = ricci_down[i, j]

print("===========================Ricci tensor=========================")
for i in range(dim):
    for j in range(i+1):
        if(ricci_down[i, j] != 0):
            print(r'R_{'+str(coor_names[i])+str(coor_names[j])+r'}=')
            print_tex(ricci_down[i, j])

for i in range(dim):
    for j in range(dim):
        for k in range(dim):
            for l in range(k+1):
                s = Riemann_tensor_up(i, j, k, l)
                if( s != 0):
                    print('R^{'+str(coor_names[i])+'}_{'+str(coor_names[j])+str(coor_names[k])+str(coor_names[l])+'}=')
                    print_tex(s)

for i in range(dim):
    for j in range(dim):
        for k in range(j+1):
            if(gam_up[i, j, k] != 0):
                print(r'\Gamma^{'+str(coor_names[i])+r'}_{'+str(coor_names[j])+str(coor_names[k])+r'}=')
                print_tex(gam_up[i, j, k])