#!/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

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

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))
    print dsolve(e, [lam(r)])

def eq2():
    r = Symbol("r")
    e = Rmn.dd(1,1)
    C = Symbol("CC")
    e = e.subs(nu(r), -lam(r))
    print dsolve(e, [lam(r)])

def eq3():
    r = Symbol("r")
    e = Rmn.dd(2,2)
    e = e.subs(nu(r), -lam(r))
    print dsolve(e, [lam(r)])

def eq4():
    r = Symbol("r")
    e = Rmn.dd(3,3)
    e = e.subs(nu(r), -lam(r))
    print dsolve(e, [lam(r)])



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 the Einstein's equations:"
    e = e.subs(nu(r), -lam(r))
    l =  dsolve(e, [lam(r)])
    pprint( Eq(lam(r), l) )
    metric = gdd.subs(lam(r), l).subs(nu(r),-l)#.combine()
    print "metric:"
    pprint( metric )

if __name__ == "__main__":
    main()
