#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))
#R^i_{ j}
ricci_mixed = sym.MutableDenseNDimArray(range(dim**2), shape=(dim, dim))
#R_{ij}
ricci_up = sym.MutableDenseNDimArray(range(dim**2), shape=(dim, dim))
#G_{ij}
einstein_down = sym.MutableDenseNDimArray(range(dim**2), shape=(dim, dim))
#G^i_{ j}
einstein_mixed = sym.MutableDenseNDimArray(range(dim**2), shape=(dim, dim))
#G^{ij}
einstein_up = sym.MutableDenseNDimArray(range(dim**2), shape=(dim, dim))


#define the coordinates u[0],u[1],...
u = sym.symarray('u',dim)
#################USER SECTION ########################
###Please define the coordinate names and metric here##########
coor_names = sym.symbols(r't, r, \theta, \phi')  #the names of the coordinates
#--------------define the metric -----------------
a = sym.symbols('a', cls=sym.Function)  #scale factor
k = sym.symbols('k')  #curvature parameter
gdown = sym.Matrix([[1, 0, 0, 0],[0, -a(u[0])**2/(1-k*u[1]**2), 0, 0],[0, 0, -a(u[0])**2*u[1]**2, 0],[0, 0, 0, -a(u[0])**2*u[1]**2*sym.sin(u[2])**2]])
#######################################################################################

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() 


## connection
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

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)

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]

##Riemann tensor and 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]

for i in range(dim):
    for j in range(dim):
        ricci_mixed[i, j] = 0
        for k in range(dim):
            ricci_mixed[i, j] += gup[i, k] * ricci_down[k, j]
        ricci_mixed[i, j] = sym.simplify(ricci_mixed[i, j])

for i in range(dim):
    for j in range(i+1):
        ricci_up[i, j] = 0
        for k in range(dim):
            ricci_up[i, j] += ricci_mixed[i, k] * gup[k, j]
        ricci_up[i, j] = sym.simplify(ricci_up[i, j])
        if(i != j):
            ricci_up[j, i] = ricci_up[i, j]

            
##===========================compute Ricci scalar ====================
ricci_scalar = 0
for k in range(dim):
    for l in range(k+1):
        if(gup[k, l] != 0):
            if(k==l):
                ricci_scalar += ricci_down[k, l]*gup[k,l]
            else:
                ricci_scalar += 2 * ricci_down[k, l]*gup[k,l]    
ricci_scalar =  sym.simplify(ricci_scalar)
##-------------------now Ricci scalar is saved-----------------------
                



##=======================compute Einstein tensor ==================
for i in range(dim):
    for j in range(i+1):
        einstein_down[i,j] = sym.simplify(ricci_down[i, j] - ricci_scalar/2 * gdown[i, j])
        einstein_up[i,j] = sym.simplify(ricci_up[i, j] - ricci_scalar/2 * gup[i, j])
    if(i != j):
        einstein_down[j, i] = einstein_down[i, j]
        einstein_up[j, i] = einstein_up[i, j]        
    
for i in range(dim):
    for j in range(dim):
        if( i == j ):
            einstein_mixed[i,j] = sym.simplify(ricci_mixed[i, j] - ricci_scalar/2)
        else:
            einstein_mixed[i,j] = ricci_mixed[i, j]



################################ OUTPUTS ###############################################
print(r"Ricci scalar = ")
print_tex(ricci_scalar)


if(input("print connections: (y/n)")=="y"):
    print("\n==========================connection========================")
    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])

if(input("print Riemann tensors: (y/n)") == "y"):                    
    print("\n==========================Riemann tensor====================")
    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)

if(input("print Ricci tensor: (y/n)")== "y"):
    print("\n===========================Ricci tensor down=========================")
    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])


    print("\n===========================Ricci tensor mixed=========================")
    for i in range(dim):
        for j in range(i+1):
            if(ricci_down[i, j] != 0):
                print(r'R^{'+str(coor_names[i])+ r'}_{\ \,' +str(coor_names[j])+r'} =')
                print_tex(ricci_mixed[i, j])
            

    print("\n===========================Ricci tensor up=========================")
    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_up[i, j])
            

if(input("print Einstein tensor: (y/n)")== "y"):                
    print("\n===========================Einstein tensor down =========================")
    for i in range(dim):
        for j in range(i+1):
            if(einstein_down[i, j] != 0):
                print(r'G_{'+str(coor_names[i])+' '+str(coor_names[j])+r'} =')
                print_tex(einstein_down[i, j])


    print("\n===========================Einstein tensor mixed =========================")
    for i in range(dim):
        for j in range(i+1):
            if(einstein_down[i, j] != 0):
                print(r'G^{'+str(coor_names[i])+ r'}_{\ \,' +str(coor_names[j])+r'} =')
                print_tex(einstein_mixed[i, j])



    print("\n===========================Einstein tensor up =========================")
    for i in range(dim):
        for j in range(i+1):
            if(einstein_down[i, j] != 0):
                print(r'G^{'+str(coor_names[i])+' '+str(coor_names[j])+r'} =')
                print_tex(einstein_up[i, j])