#sample script: calculate Einstein tensor
#----------by Zhiqi Huang-------------
import sympy as sym
dim = 4
numcoor = sym.symarray('x',dim) #u is a reserved symbol
#myprinter = sym.pretty   ## preview on the screen
myprinter = sym.latex  ##generate a latex file


##FRW non-flat metric, k is the spatial curvature parameter
t, r, theta, phi = sym.symbols(r't, r, \theta, \phi')
coors = [t, r, theta, phi]
a = sym.symbols('a', cls=sym.Function)
k = sym.symbols('k')
metric = sym.diag( 1, -a(t)**2/(1-k*r**2), -a(t)**2*r**2, -a(t)**2 * r**2 * sym.sin(theta)**2 )

print_connection = True
print_ricci_tensor = True
print_ricci_scalar = True
print_einstein_tensor = True



#================= DO NOT CHANGE CODES BELOW ========================
coors_to_num = [ (coors[i], numcoor[i]) for i in range(dim)]
num_to_coors = [ (numcoor[i], coors[i]) for i in range(dim)]
gdown = metric.subs( coors_to_num )

if(myprinter == sym.latex):
    print(r'\documentclass[12pt]{article}\begin{document}')
    print(r'\section{Metric}')
    print(r'\begin{equation} ds^2=')
    counter = 0    
    for i in range(dim):
        for j in range(i+1):
            if(metric[i, j] != 0):
                counter += 1
                if(i == j):
                    mstr = sym.latex(metric[i, j])
                    cstr = r'd{' +  str(coors[i]) + '}^2'
                else:
                    mstr = sym.latex(sym.simplify(2*metric[i, j]))
                    cstr = r'd{' + str(coors[i]) + r'} d{' + str(coors[j]) + r'}'
                if(mstr == '1'):
                    if(counter > 1):
                        mstr = '+'
                    else:
                        mstr = ''
                elif(mstr == '-1'):
                    mstr = '-'
                elif(mstr[0] != '-' and counter > 1):
                    mstr = '+' + mstr
                print(mstr + cstr)
    print(r'\end{equation}')            

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

def do_print_connection():
    if(myprinter == sym.latex):
        print("\section{Connections}")        
        print(r'\begin{eqnarray}')
    else:
        print("%--------------connections----------")        
    counter = 0
    for i in range(dim):
        for j in range(dim):
            for k in range(j+1):
                if(gam_up[i, j, k] != 0):
                    counter += 1
                    if(myprinter == sym.latex):
                        if(counter > 1):
                            print(r' \\')
                        print(r'\Gamma^{', coors[i],  r'}_{', coors[j], ' ', coors[k],  r'} &=& ', myprinter(gam_up[i, j, k].subs(num_to_coors)))
                    else:
                        print(r'Gamma^', coors[i],  r'_', coors[j], ' ', coors[k],  r' = ')
                        print(myprinter(gam_up[i, j, k].subs(num_to_coors)) )
                        print("\n")
    if(myprinter == sym.latex):
        print(r'\end{eqnarray}')

def do_print_ricci_tensor():    
    if(myprinter == sym.latex):
        print(r'\section{Ricci Tensor $R_{\mu\nu}$}')                
        print(r'\begin{eqnarray}')
    else:
        print(r'%%%%%%%%%%%%%%%%%%% Ricci tensor R_{ij} %%%%%%%%%%%%%%')
    counter = 0        
    for i in range(dim):
        for j in range(i+1):
            if(ricci_down[i, j] != 0):
                counter += 1
                if(myprinter == sym.latex):
                    if(counter > 1):
                        print( r' \\ ' )
                    print(r'R_{', coors[i], r' ', coors[j], r'} &=& ', myprinter(ricci_down[i, j].subs(num_to_coors)))
                else:
                    print(r'R_', coors[i],  ' ' , coors[j], r' = ' )
                    print(myprinter(ricci_down[i, j].subs(num_to_coors)))
                    print("\n")
    if(myprinter == sym.latex):
        print(r'\end{eqnarray}')                                
        print(r'\section{Ricci Tensor $R^{\mu\nu}$}')                
        print(r'\begin{eqnarray}')
    else:
        print(r'%%%%%%%%%%%%%%% Ricci tensor R^{ij} %%%%%%%%%%%%%%%%%%')
    counter = 0
    for i in range(dim):
        for j in range(i+1):
            if(ricci_up[i, j] != 0):
                counter += 1
                if(myprinter == sym.latex):
                    if(counter > 1):
                        print(r' \\')                    
                    print(r'R^{', coors[i], r' ', coors[j], r'} &=& ' + myprinter(ricci_up[i, j].subs(num_to_coors)))
                else:
                    print(r'R^', coors[i],  ' ' , coors[j], r' = ' )
                    print(myprinter(ricci_up[i, j].subs(num_to_coors)))
                    print("\n")                    
    if(myprinter == sym.latex):
        print(r'\end{eqnarray}')                                
        print(r'\section{Ricci Tensor $R^\mu_{\ \nu}$}')                
        print(r'\begin{eqnarray}')
    else:
        print(r'%%%%%%%%%%%%%% Ricci tensor R^i_j %%%%%%%%%%%%%%%%%%%%')
    counter  = 0
    for i in range(dim):
        for j in range(i+1):
            if(ricci_mixed[i, j] != 0):
                counter += 1
                if(myprinter == sym.latex):
                    if(counter > 1):
                        print(r' \\')
                    print(r'R^{', coors[i], r'}_{\ ', coors[j], r'} &=& ' + myprinter(ricci_mixed[i, j].subs(num_to_coors)) )                    
                else:
                    print(r'R^', coors[i],  r'_' , coors[j], r' = ' )
                    print(myprinter(ricci_mixed[i, j].subs(num_to_coors)))
                    print("\n")                    
    if(myprinter == sym.latex):
        print(r'\end{eqnarray}')                

        
def do_print_ricci_scalar():
    if(myprinter == sym.latex):
        print(r'\section{Ricci scalar}')
        print(r'\begin{equation}')
        print('R = ', myprinter(ricci_scalar.subs(num_to_coors)))
        print(r'\end{equation}')        
    else:
        print("%%%%%%%%%%%%% Ricci scalar %%%%%%%%%%%%%%%")
        print(ricci_scalar.subs(num_to_coors))
        print("\n")    

def do_print_einstein_tensor():
    if(myprinter == sym.latex):
        print(r'\section{Einstein Tensor $G_{\mu\nu}$}')                
        print(r'\begin{eqnarray}')
    else:
        print(r'%%%%%%%%%%%%%%%%%%% Einstein tensor G_{ij} %%%%%%%%%%%%%%')
    counter = 0        
    for i in range(dim):
        for j in range(i+1):
            if(einstein_down[i, j] != 0):
                counter += 1
                if(myprinter == sym.latex):
                    if(counter > 1):
                        print( r' \\ ' )
                    print(r'G_{', coors[i], r' ', coors[j], r'} &=& ', myprinter(einstein_down[i, j].subs(num_to_coors)))
                else:
                    print(r'G_', coors[i],  ' ' , coors[j], r' = ' )
                    print(myprinter(einstein_down[i, j].subs(num_to_coors)))
                    print("\n")    
    if(myprinter == sym.latex):
        print(r'\end{eqnarray}')                                
        print(r'\section{Einstein Tensor $G^{\mu\nu}$}')                
        print(r'\begin{eqnarray}')
    else:
        print(r'%%%%%%%%%%%%%%% Einstein tensor G^{ij} %%%%%%%%%%%%%%%%%%')
    counter = 0
    for i in range(dim):
        for j in range(i+1):
            if(einstein_up[i, j] != 0):
                counter += 1
                if(myprinter == sym.latex):
                    if(counter > 1):
                        print(r' \\')                    
                    print(r'G^{', coors[i], r' ', coors[j], r'} &=& ' + myprinter(einstein_up[i, j].subs(num_to_coors)))
                else:
                    print(r'G^', coors[i],  ' ' , coors[j], r' = ' )
                    print(myprinter(einstein_up[i, j].subs(num_to_coors)))
                    print("\n")                        
    if(myprinter == sym.latex):
        print(r'\end{eqnarray}')                                
        print(r'\section{Einstein Tensor $G^\mu_{\ \nu}$}')                
        print(r'\begin{eqnarray}')
    else:
        print(r'%%%%%%%%%%%%%% Einstein tensor G^i_j %%%%%%%%%%%%%%%%%%%%')
    counter  = 0
    for i in range(dim):
        for j in range(i+1):
            if(einstein_mixed[i, j] != 0):
                counter += 1
                if(myprinter == sym.latex):
                    if(counter > 1):
                        print(r' \\')
                    print(r'G^{', coors[i], r'}_{\ ', coors[j], r'} &=& ' + myprinter(einstein_mixed[i, j].subs(num_to_coors)) )                    
                else:
                    print(r'G^', coors[i],  r'_' , coors[j], r' = ' )
                    print(myprinter(einstein_mixed[i, j].subs(num_to_coors)))
                    print("\n")                        
    if(myprinter == sym.latex):
        print(r'\end{eqnarray}')                
            

gup = gdown ** -1

detg = gdown.det()

def connection_down(i, j, k):
    return (sym.diff(gdown[i, j], numcoor[k]) + sym.diff(gdown[i, k], numcoor[j]) - sym.diff(gdown[j, k], numcoor[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]

##now we have both gam_down and gam_up saved
if(print_connection):
    do_print_connection()
##===================compute Ricci tensor ==============================
if (not print_ricci_tensor and not print_ricci_scalar and not print_einstein_tensor):
    if(myprinter == sym.latex):
        print(r'\end{document}')
    exit()

def Riemann_tensor_up(i, j, k, l): ## R^i_{  jkl}
    R = sym.diff(gam_up[i, j, k], numcoor[l]) - sym.diff(gam_up[i, j, l], numcoor[k])
    for m in range(dim):
        R += gam_up[i, m, l] * gam_up[m, j, k] - gam_up[i, m, k] * gam_up[m, j, l]
    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, j, k)
        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]

if do_print_ricci_tensor:
    do_print_ricci_tensor()
## ------------------ now Ricci tensor is saved-----------------------

##===========================compute Ricci scalar ====================
if(print_ricci_scalar or print_einstein_tensor):
    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)
if(print_ricci_scalar):
    do_print_ricci_scalar()
##-------------------now Ricci scalar is saved-----------------------

##=======================compute Einstein tensor ==================
if print_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]
    do_print_einstein_tensor()

if(myprinter == sym.latex):
    print(r'\end{document}')