#sample code: calculate Ricci tensor for spherical metric
#----------by Zhiqi Huang--for the course ''General Relativity''-------------

import sympy as sym
#optmize printing
sym.init_printing()

#dimension of the spacetime
dim = 4

##allocate space to save connections, Riemann tensor, and Ricci tensor
gam_down = sym.MutableDenseNDimArray(range(dim**3), shape=(dim, dim, dim))
gam_up = sym.MutableDenseNDimArray(range(dim**3), shape=(dim, dim, dim))
ricci_down = sym.MutableDenseNDimArray(range(dim**2), shape=(dim, dim))


#define u[0], u[1], ... in a single line
u = sym.symarray('u',dim)

#for better output
t, r, theta, phi = sym.symbols(r't, r, theta, phi')

#define the Newtonian potential and spatial curvature functions
Phi = sym.Function('Phi')
Psi = sym.Function('Psi')

#covariant metric
gdown = sym.diag( sym.exp(2*Phi(u[0], u[1])), -sym.exp(-2*Psi(u[0], u[1])) , -u[1]**2, -u[1]**2 * sym.sin(u[2]) **2 )

#contravariant metric
gup = gdown ** -1

#determinant of the covariant metric
detg = gdown.det()

#\Gamma_{ijk}$
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]

#\Gamma^i_{\ jk}
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
                
def Riemann_tensor_up(i, j, k, l): ## R^i_{  jkl}
    R = sym.diff(gam_up[i, j, k], u[l]) - sym.diff(gam_up[i, j, l], u[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]
                                    
                                       
# now Ricci tensor is saved


#print latex
print(r'\begin{eqnarray}')
for i in range(dim):
    for j in range(i+1):
        Rij = ricci_down[i, j].subs(u[0], t).subs(u[1], r).subs(u[2], theta).subs(u[3], phi)
        if(Rij != 0):
            print(str(r'R_{'+str(i)+str(j)+r'} &=& '+sym.latex(Rij) +r' \nonumber \\').replace(r'\Phi{\left (t,r \right )}', r'\Phi').replace(r'\Psi{\left (t,r \right )}', r'\Psi').replace(r'\frac{\partial}{\partial t} \Phi', r'\Phi_{,t}').replace(r'\frac{\partial}{\partial t} \Psi', r'\Psi_{,t}').replace(r'\frac{\partial}{\partial r} \Phi', r'\Phi_{,r}').replace(r'\frac{\partial}{\partial r} \Psi', r'\Psi_{, r}').replace(r'\frac{\partial^{2}}{\partial r^{2}}  \Phi', r'\Phi_{,r,r}').replace(r'\frac{\partial^{2}}{\partial t^{2}}  \Psi', r'\Phi_{,t,t}') )
print(r'\end{eqnarray}')