# Siehe VARTC 2026/27 (siehe http://itp.uni-frankfurt.de/~hanauske/VARTC/index.html , Vorlesung 6)
# Python Skript zum Visualisieren der Ausgabe des C++ Programms GeodesicSchwarzschild.cpp + GeodesicSchwarzschild.hpp
# C++ Programm zum Loesen der Geodaetengleichung in vorgegebener Schwarzschild-Raumzeit (nicht-rotierendes, ungeladenes Schwarzes Loch)
# Ausgabe-Datei des C++ Programms: geodesic.dat

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import rcParams
import matplotlib.gridspec as gridspec
import matplotlib.animation as animation
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm

# Globales Setup für die Darstellung
rcParams.update({
    'figure.figsize' : [12, 6],
    'axes.titlesize' : 14,
    'axes.labelsize' : 14,
    'xtick.labelsize': 12,
    'ytick.labelsize': 12
})

# Daten laden
M = 1
data = np.genfromtxt("./geodesic.dat")
N = len(data)

# Berechnung des Drehimpulses l
l = data[0, 3]**2 * data[0, 7]

def VeffRez(r):
    # np.real_if_close fängt minimale numerische Imaginärteile ab
    return np.real_if_close(np.sqrt(1 - 2*M/r) * np.sqrt(1 + l**2/r**2))

# Gitter für die 3D-Fläche
R = np.linspace(2.00001, 20, 100)
PHI = np.linspace(0, 2*np.pi, 100)
R, PHI = np.meshgrid(R, PHI)
X = R * np.cos(PHI)
Y = R * np.sin(PHI)
ZF = np.real(np.sqrt(8*M*(R - 2*M)))

# Trajektorien-Daten aus der Datei
r_data = data[:, 3]
r_max = np.max(r_data)
r_min = np.min(r_data)
phi_data = data[:, 4]
x_all = r_data * np.cos(phi_data)
y_all = r_data * np.sin(phi_data)
z_all = np.real(np.sqrt(8*M * (r_data - 2*M)))
z_shadow = np.zeros_like(z_all)

# Frame-Indizes
anzframes = 100 # Anzahl der Bilder im Film
frame_indices = np.linspace(0, N - 1, anzframes, dtype=int)

# Figure-Setup
fig = plt.figure()
gs = gridspec.GridSpec(1, 2, width_ratios=[1.8, 1], wspace=0.25)
ax1 = plt.subplot(gs[0], projection='3d')
ax2 = plt.subplot(gs[1])

# Achsen-Design (Transparente Wände)
for axis in [ax1.xaxis, ax1.yaxis, ax1.zaxis]:
    axis.pane.fill = False
    axis.pane.set_edgecolor('white')

xylim = 18
zlim = 15
ax1.set_xlabel(r"$\rm x \,[km]$")
ax1.set_ylabel(r"$\rm y \,[km]$")
ax1.set_zlabel("Einbettung z")
ax1.set_xlim(-xylim, xylim)
ax1.set_ylim(-xylim, xylim)
ax1.set_zlim(0, zlim)
ax1.view_init(azim=35, elev=20)
ax1.grid(True, linestyle=':', alpha=0.5)

# Statische 3D-Fläche
ax1.plot_surface(X, Y, ZF, cmap=cm.nipy_spectral, linewidth=0, alpha=0.3, rstride=2, cstride=2)
ax1.contourf(X, Y, ZF, levels=30, zdir='z', offset=0, cmap=cm.nipy_spectral, alpha=0.2)

# Rechtes Potenzialdiagramm
rval = np.linspace(2.85, 28, 500)
ax2.plot(rval, VeffRez(rval), c="black", linewidth=1.5, label=r"$V_{\rm eff}(r)$")
ax2.plot([r_min,r_max], [VeffRez(r_min),VeffRez(r_min)], c="blue", linewidth=0.5, linestyle=':')
ax2.set_xlabel(r"$\rm Radius \, r \,[km]$")
ax2.set_ylabel(r"$\rm V(r)$")
ax2.set_xlim(2.85, 28)
ax2.grid(True, linestyle=':', alpha=0.5)

# Dynamische Elemente initialisieren
line_3d, = ax1.plot([], [], [], c="blue", linewidth=1.2, label="Trajektorie")
line_shadow, = ax1.plot([], [], [], c="gray", linewidth=1, linestyle='--')
scatter_3d = ax1.scatter([], [], [], c="blue", s=50)
scatter_shadow = ax1.scatter([], [], [], c="gray", s=30)
scatter_pot = ax2.scatter([], [], c="blue", s=60, zorder=5)

def init():
    line_3d.set_data([], [])
    line_3d.set_3d_properties([])
    line_shadow.set_data([], [])
    line_shadow.set_3d_properties([])
    scatter_3d.set_offsets(np.empty((0, 2)))
    scatter_shadow.set_offsets(np.empty((0, 2)))
    scatter_pot.set_offsets(np.empty((0, 2)))
    return line_3d, line_shadow, scatter_3d, scatter_shadow, scatter_pot

def animate(i):
    idx = frame_indices[i]

    # Linien (Vergangenheit der Trajektorie)
    line_3d.set_data(x_all[:idx+1], y_all[:idx+1])
    line_3d.set_3d_properties(z_all[:idx+1])

    line_shadow.set_data(x_all[:idx+1], y_all[:idx+1])
    line_shadow.set_3d_properties(z_shadow[:idx+1])

    # Aktuelle Positionen (Punkte)
    scatter_3d.set_offsets(np.c_[x_all[idx], y_all[idx]])
    scatter_3d.set_3d_properties([z_all[idx]], zdir='z')

    scatter_shadow.set_offsets(np.c_[x_all[idx], y_all[idx]])
    scatter_shadow.set_3d_properties([0], zdir='z')

    # Potenzialverlauf
    current_r = r_data[idx]
    scatter_pot.set_offsets(np.c_[current_r, VeffRez(current_r)])

    return line_3d, line_shadow, scatter_3d, scatter_shadow, scatter_pot

# Animation erstellen und speichern
ani = animation.FuncAnimation(
    fig, animate, init_func=init, frames=anzframes, interval=80, blit=False
)

plt.close(fig)
ani.save('GeodesicSchwarzschild.mp4', writer='ffmpeg', fps=15, dpi=150)