#sample script: calculate connection, Ricci scalar and Einstein tensor in Newtonian gauge
#----------by Zhiqi Huang-----------huangzhq25@mail.sysu.edu.cn ----

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


##Newtonian Gauge metric
t, x, y, z, epsilon = sym.symbols('tau, x, y, z, epsilon')
coors = [t, x, y, z]
a, Phi, Psi = sym.symbols('a, Phi, Psi', cls=sym.Function)
metric = sym.diag( a(t)**2*(1+2*Phi(t, x, y, z)*epsilon), - a(t)**2*(1-2*Psi(t, x, y, z)*epsilon), - a(t)**2*(1-2*Psi(t, x, y, z)*epsilon), - a(t)**2*(1-2*Psi(t, x, y, z)*epsilon) )

print_connection = True
print_ricci_scalar = True
print_einstein_tensor = True  #only print mixed G^mu_nu


def first_order( expr ):
    return sym.series(expr , epsilon, 0, 2) 


def first_order_2D( expr ):
    for i in range(dim):
        for j in range(dim):
            expr[i, j] = first_order(expr[i, j])

def print_compact_latex( mystr):
    print(mystr.replace(r'\Phi{\left(\tau,x,y,z \right)}',r'\Phi ').replace(r'\Psi{\left(\tau,x,y,z \right)}',r'\Psi ').replace(r'a{\left(\tau \right)}',r'a ').replace(r'+ O\left(\epsilon^{2}\right)', '').replace(r'a^{2}{\left(\tau \right)}', r'a^2 ').replace(r'a^{3}{\left(\tau \right)}', r'a^3 ').replace(r'a^{4}{\left(\tau \right)}', r'a^4 ').replace(r'\epsilon', '').replace(r'\frac{d}{d \tau} a ', r"{a'} ").replace(r'\frac{d^{2}}{d \tau^{2}} a ', r"{a''}").replace(r'\frac{\partial^{2}}{\partial x^{2}}', r'\partial_{xx}').replace(r'\frac{\partial^{2}}{\partial y^{2}}', r'\partial_{yy}').replace(r'\frac{\partial^{2}}{\partial z^{2}}', r'\partial_{zz}').replace(r'\frac{\partial}{\partial \tau}', r'{\partial_\tau}').replace(r'\frac{\partial^{2}}{\partial \tau^{2}}', r'{\partial_{\tau\tau}}').replace(r'\frac{\partial}{\partial x}',r'{\partial_x}').replace(r'\frac{\partial}{\partial y}',r'{\partial_y}').replace(r'\frac{\partial}{\partial z}',r'{\partial_z}').replace(r'\frac{\partial^{2}}{\partial y\partial z}',r'{\partial_{yz}}').replace(r'\frac{\partial^{2}}{\partial z\partial y}',r'{\partial_{zy}}').replace(r'\frac{\partial^{2}}{\partial x\partial y}',r'{\partial_{xy}}').replace(r'\frac{\partial^{2}}{\partial y\partial x}',r'{\partial_{yx}}').replace(r'\frac{\partial^{2}}{\partial x\partial z}',r'{\partial_{xz}}').replace(r'\frac{\partial^{2}}{\partial z\partial x}',r'{\partial_{zx}}').replace(r'\frac{\partial^{2}}{\partial x\partial \tau}',r'{\partial_{x\tau}}').replace(r'\frac{\partial^{2}}{\partial y\partial \tau}',r'{\partial_{y\tau}}').replace(r'\frac{\partial^{2}}{\partial z\partial \tau}',r'{\partial_{z\tau}}'))

#================= 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[10pt]{article}\begin{document}')
    print(r'\section{Metric}')
    print(r'$$ ds^2 = a^2(\tau)\left[(1+2\Phi(\tau, x, y,z))d\tau^2 - (1-2\Psi(\tau, x, y,z))(dx^2 + dy^2+dz^2)\right]$$')        

##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' \\')
                        mystr = r'\Gamma^{' + myprinter(coors[i])+  r'}_{' + myprinter(coors[j])+ ' ' + myprinter(coors[k])+  r'} &=& ' + myprinter(gam_up[i, j, k].subs(num_to_coors))
                        print_compact_latex(mystr)
                    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_scalar():
    if(myprinter == sym.latex):
        print(r'\section{Ricci scalar}')
        print(r'{\tiny \begin{equation}')
        print_compact_latex('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 Tensors}')                
        print(r' {\tiny \begin{eqnarray}')
    else:
        print(r'%%%%%%%%%%%%%% Einstein tensor G^i_j %%%%%%%%%%%%%%%%%%%%')
    counter  = 0
    for i in range(dim):
        for j in range(dim):
            if(einstein_mixed[i, j] != 0):
                counter += 1
                if(myprinter == sym.latex):
                    if(counter > 1):
                        print(r' \\')
                    mystr = r'G^{' + myprinter(coors[i]) +  r'}_{\ ' + myprinter(coors[j]) + r'} &=& ' + myprinter(einstein_mixed[i, j].subs(num_to_coors))
                    print_compact_latex(mystr)
                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

first_order_2D(gup)

detg = first_order(gdown.det())


def connection_down(i, j, k):
    return first_order((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 first_order(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_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 first_order(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] = first_order(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] = first_order(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] = first_order(sym.simplify(ricci_up[i, j]))
        if(i != j):
            ricci_up[j, i] = ricci_up[i, j]

## ------------------ 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 =  first_order(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] = first_order(ricci_down[i, j] - ricci_scalar/2 * gdown[i, j])
            einstein_up[i,j] = first_order(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] = first_order(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}')