#!/usr/bin/env python
 
"""
This example calculates the Ricci tensor from the metric and does this
on the example of Schwarzschild solution.
 
If you want to derive this by hand, follow the wiki page here:
 
http://en.wikipedia.org/wiki/Deriving_the_Schwarzschild_solution
 
Also read the above wiki and follow the references from there if something is
not clear, like what the Ricci tensor is, etc.
 
"""
 
from sympy import exp, Symbol, sin, Rational, Derivative, dsolve, Function, \
                  Matrix, Eq, pprint, Pow, classify_ode
 
def grad(f,X):
    a=[]
    for x in X:
        a.append(f.diff(x))
    return a
 
def d(m,x):
    return grad(m[0,0],x)
 
class MT(object):
    def __init__(self,m):
        self.gdd=m
        self.guu=m.inv()
 
    def __str__(self):
        return "g_dd =\n" + str(self.gdd)
 
    def dd(self,i,j):
        return self.gdd[i,j]
 
    def uu(self,i,j):
        return self.guu[i,j]
 
class G(object):
    def __init__(self,g,x):
        self.g = g
        self.x = x
 
    def udd(self,i,k,l):
        g=self.g
        x=self.x
        r=0
        for m in [0,1,2,3]:
            r+=g.uu(i,m)/2 * (g.dd(m,k).diff(x[l])+g.dd(m,l).diff(x[k]) \
                    - g.dd(k,l).diff(x[m]))
        return r
 
class Riemann(object):
    def __init__(self,G,x):
        self.G = G
        self.x = x
 
    def uddd(self,rho,sigma,mu,nu):
        G=self.G
        x=self.x
        r=G.udd(rho,nu,sigma).diff(x[mu])-G.udd(rho,mu,sigma).diff(x[nu])
        for lam in [0,1,2,3]:
            r+=G.udd(rho,mu,lam)*G.udd(lam,nu,sigma) \
                -G.udd(rho,nu,lam)*G.udd(lam,mu,sigma)
        return r
 
class Ricci(object):
    def __init__(self,R,x):
        self.R = R
        self.x = x
        self.g = R.G.g
 
    def dd(self,mu,nu):
        R=self.R
        x=self.x
        r=0
        for lam in [0,1,2,3]:
            r+=R.uddd(lam,mu,lam,nu)
        return r
 
    def ud(self,mu,nu):
        r=0
        for lam in [0,1,2,3]:
            r+=self.g.uu(mu,lam)*self.dd(lam,nu)
        return r.expand()
 
def curvature(Rmn):
    return Rmn.ud(0,0)+Rmn.ud(1,1)+Rmn.ud(2,2)+Rmn.ud(3,3)
 
#class nu(Function):
#    def getname(self):
#        return r"\nu"
#        return r"nu"
 
#class lam(Function):
#    def getname(self):
#        return r"\lambda"
#        return r"lambda"
nu = Function("nu")
lam = Function("lambda")
 
t=Symbol("t")
r=Symbol("r")
theta=Symbol(r"theta")
phi=Symbol(r"phi")
 
#general, spherically symmetric metric
gdd=Matrix((
    (-exp(nu(r)),0,0,0),
    (0, exp(lam(r)), 0, 0),
    (0, 0, r**2, 0),
    (0, 0, 0, r**2*sin(theta)**2)
    ))
#spherical - flat
#gdd=Matrix((
#    (-1, 0, 0, 0),
#    (0, 1, 0, 0),
#    (0, 0, r**2, 0),
#    (0, 0, 0, r**2*sin(theta)**2)
#    ))
#polar - flat
#gdd=Matrix((
#    (-1, 0, 0, 0),
#    (0, 1, 0, 0),
#    (0, 0, 1, 0),
#    (0, 0, 0, r**2)
#    ))
#polar - on the sphere, on the north pole
#gdd=Matrix((
#    (-1, 0, 0, 0),
#    (0, 1, 0, 0),
#    (0, 0, r**2*sin(theta)**2, 0),
#    (0, 0, 0, r**2)
#    ))
g=MT(gdd)
X=(t,r,theta,phi)
Gamma=G(g,X)
Rmn=Ricci(Riemann(Gamma,X),X)
 
def pprint_Gamma_udd(i,k,l):
    pprint(Eq(Symbol('Gamma^%i_%i%i' % (i,k,l)), Gamma.udd(i,k,l)))
 
def pprint_Rmn_dd(i,j):
    pprint(Eq(Symbol('R_%i%i' % (i,j)), Rmn.dd(i,j)))
 
 
# from Differential Equations example
def eq1():
    r = Symbol("r")
    e = Rmn.dd(0,0)
    e = e.subs(nu(r), -lam(r))
    pprint(dsolve(e, lam(r)))
 
def eq2():
    r = Symbol("r")
    e = Rmn.dd(1,1)
    C = Symbol("CC")
    e = e.subs(nu(r), -lam(r))
    pprint(dsolve(e, lam(r)))
 
def eq3():
    r = Symbol("r")
    e = Rmn.dd(2,2)
    e = e.subs(nu(r), -lam(r))
    pprint(dsolve(e, lam(r)))
 
def eq4():
    r = Symbol("r")
    e = Rmn.dd(3,3)
    e = e.subs(nu(r), -lam(r))
    pprint(dsolve(e, lam(r)))
    pprint(dsolve(e, lam(r), 'best'))
 
 
 
def main():
 
    print "Initial metric:"
    pprint(gdd)
    print "-"*40
    print "Christoffel symbols:"
    pprint_Gamma_udd(0,1,0)
    pprint_Gamma_udd(0,0,1)
    print
    pprint_Gamma_udd(1,0,0)
    pprint_Gamma_udd(1,1,1)
    pprint_Gamma_udd(1,2,2)
    pprint_Gamma_udd(1,3,3)
    print
    pprint_Gamma_udd(2,2,1)
    pprint_Gamma_udd(2,1,2)
    pprint_Gamma_udd(2,3,3)
    print
    pprint_Gamma_udd(3,2,3)
    pprint_Gamma_udd(3,3,2)
    pprint_Gamma_udd(3,1,3)
    pprint_Gamma_udd(3,3,1)
    print"-"*40
    print"Ricci tensor:"
    pprint_Rmn_dd(0,0)
    e =  Rmn.dd(1,1)
    pprint_Rmn_dd(1,1)
    pprint_Rmn_dd(2,2)
    pprint_Rmn_dd(3,3)
    #print
    #print "scalar curvature:"
    #print curvature(Rmn)
    print "-"*40
    print "Solve Einstein's equations:"
    e = e.subs(nu(r), -lam(r))
    l =  dsolve(e, lam(r))
    pprint(l)
    metric = gdd.subs(lam(r), l).subs(nu(r),-l)#.combine()
    print "metric:"
    pprint(metric)
 
if __name__ == "__main__":
    main()