# Version vom 15.5.23
# Geodäten im H^2 (Poincaré disk model) -- https://en.wikipedia.org/wiki/Poincar%C3%A9_disk_model

from sympy import symbols, diag
#https://de.wikipedia.org/wiki/SymPy
#https://www.sympy.org/en/index.html 

from einsteinpy.symbolic import MetricTensor, RicciTensor, RicciScalar, ChristoffelSymbols, RiemannCurvatureTensor
#https://einsteinpy.org/

import numpy as np
#https://numpy.org/

from scipy.integrate import odeint
#https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.odeint.html

import matplotlib.pyplot as plt
#https://matplotlib.org/

#Geodätengleichung (2d) als System 2.Ordnung
#ch Christoffelsymbole
def geod(y, t, ch):
    theta1, phi1, tv, pv = y
# ch ist symbolisch, subs - substituiert die Variable theta mit theta1 (was im Aufruf dann immer eine Zahl ist) etc    
    ch2 = ch.tensor().subs([(x1, theta1),(x2, phi1)])
# theta1'=tv, phi1'=pv, tv'=- \Gamma^{theta1}_{ij} x^ix^j, pv'=- \Gamma^{phi1}_{ij} x^ix^j,  (x^0=theta1, x^1=phi1)
    dydt = [tv, pv, -ch2[0,1,0]*tv**2-2*ch2[0,0,1]*tv*pv-ch2[0,1,1]*pv**2, -ch2[1,0,0]*tv**2-2*ch2[1,0,1]*tv*pv-ch2[1,1,1]*pv**2]
    return dydt


syms = symbols('x1 x2')
x1, x2  = syms
    

# Metrik in sphärischen Koordinaten - vgl. ART-Skript Bsp II.1.13
g = MetricTensor(diag(4/(1-x1**2-x2**2)**2, 4/(1-x1**2-x2**2)**2).tolist(), syms)
print('Metrik:', g.tensor())
#print('g[1=theta,2=phi]=', g[0,1], '\n')

ch = ChristoffelSymbols.from_metric(g)
# \Gamma_{ij}^k: äußerste Klammer k, innerste Klammer i
print('Christoffelsymbole:', ch.tensor())

#Konvention Riemann-Tensor auf https://docs.einsteinpy.org/en/stable/api/symbolic/riemann.html
# Riem^i_{jkl} wird mit Riem[i,k,l,j] ausgegeben 
Riem = RiemannCurvatureTensor.from_metric(g)
print('Riem - als (1,4)-Tensor:', Riem.tensor())
print('Riem^1_{010}=', Riem[1,0,1,0], '\n')

Ric = RicciTensor.from_metric(g)
print('Ricci - als (0,2)-Tensor:', Ric.tensor())
uricci = Ric.change_config(metric=g)
print('Ricci - als (1,1)-Tensor:', uricci.tensor(), '\n')

R = RicciScalar.from_riccitensor(Ric)
R.simplify()
print('Skalarkrümmung:', R.tensor())



#############################
# Geodätengleichung (2d) -- Definition der Geodätengleichung oben
y0 = [0, 0, 0.2,0.1]
y1 = [0, 0, 0.2,0.2]
y2 = [0, 0.5, 1, -1]

t = np.linspace(0, 6,10000)

sol = odeint(geod, y0, t, args=(ch,))
sol2 = odeint(geod, y1, t, args=(ch,))

t2 = np.linspace(0, 0.8, 1000)
sol3 = odeint(geod, y2, t2, args=(ch,))


#Plotten der Lösung
fig, ax = plt.subplots()

x1 , x2 , xv, yv= zip(*sol)
ax.plot(x1, x2, color='red')

x1 , x2 , xv, yv= zip(*sol2)
ax.plot(x1, x2)

x1 , x2 , xv, yv= zip(*sol3)
ax.plot(x1, x2)


plt.xlim([-1, 1])
plt.ylim([-1, 1])

#Kreis plotten
u = np.linspace(0, 2*np.pi, 500)
ax.plot(np.cos(u), np.sin(u), color='lightgray')


plt.show()