# # Numerische Berechnung der Trajektorien des harmonischen Oszillators

# Hilfreiche Pakete
import numpy as np
import matplotlib.pyplot as plt
import math

# Konstanten
from scipy.constants import g

# Sinnvolle Konstanten
def run_sim(gamma = 0, dt = 1e-2):
    x0 = 1   
    v0 = 0
    omega0 = 1
    tmax = 20 # Max time for the algorithm to terminate

    # Initial values
    x_t = [x0]
    x_t_exact = [x0]
    v_t = [v0]
    v_t_exact = [v0]
    T_t = [0.5*v0*v0]
    V_t = [0.5*omega0*omega0*x0*x0]
    E_t = [T_t[0]+V_t[0]]

    # exakter Wert der Energie:
    E_exact = E_t[0]
    Ediff_t = [0]

    # Implementation of the Euler-Algorithm
    for i in range(int(tmax // dt)):
        f = v_t[-1]
        x = x_t[-1] + f * dt # Calc new position

        # Change this for the damped one
        #v = v_t[-1] - dt * omega0 ** 2 * x_t[-1] # undamped
        v = v_t[-1] + dt * ( -2 * gamma * v_t[-1] - omega0 ** 2 * x_t[-1]) # damped (general)

        x_t.append(x)
        v_t.append(v)

        # Q: Wofuer werden die exacten Werte gebraucht wenn sich die
        # exacte Energie nicht aendert?
        x_exact = x0 * np.exp(-gamma * i * dt) * np.cos(omega0 * i * dt)
        v_exact = -omega0 * x0 * np.exp(-gamma * i * dt) * np.sin(omega0 * i * dt)

        x_t_exact.append(x_exact)
        v_t_exact.append(v_exact)

        # calculate the energy based on current velocity and position
        T = 1/2 * v ** 2
        V = 1/2 * omega0 ** 2 * x ** 2
        E = T + V
        T_t.append(T)
        V_t.append(V)
        E_t.append(E)
        Ediff_t.append(E-E_exact)



    # TIPP: Zum erstellen mehrerer Plots auf einmal, siehe z.B.:
    # https://matplotlib.org/3.1.1/gallery/subplots_axes_and_figures/subplot.html

    kwargs = {'c':'b'}
    font_kwargs = {'fontsize':14}
    times = np.arange(0,len(x_t))*dt
    abs_max = max(x_t, key=abs)

    fig,ax =  plt.subplots(1,4,figsize=(18,7))

    #actual plots
    ax[0].plot(times,x_t,**kwargs)
    ax[0].plot(times,x_t_exact,'r--', label="exakt") # add exact values to plot
    ax[1].plot(x_t,v_t,**kwargs)
    ax[1].plot(x_t_exact,v_t_exact,'r--') # add exact values to plot
    ax[2].plot(times,T_t,label = "T")
    ax[2].plot(times,V_t,label = "V")
    ax[2].plot(times,E_t,label = "E = T+V")
    ax[3].plot(times,Ediff_t,label = "Fehler in der Energie")

    #style changes
    ax[0].set_xlim(0,len(x_t)*dt)
    ax[0].set_ylim(-abs_max,abs_max)
    ax[0].set_xlabel("t",**font_kwargs)
    ax[0].set_ylabel("$x(t)$",**font_kwargs)

    ax[1].set_xlabel("$x(t)$",**font_kwargs)
    ax[1].set_ylabel("$v(t)$",**font_kwargs)

    ax[2].set_xlim(0,len(x_t)*dt)
    ax[2].set_ylim(-abs_max,abs_max)
    ax[2].set_xlabel("t",**font_kwargs)
    ax[2].set_ylabel("Energies",**font_kwargs)

    ax[2].set_xlim(0,len(x_t)*dt)
    ax[2].set_ylim(-abs_max,abs_max)
    ax[2].set_xlabel("t",**font_kwargs)
    ax[2].set_ylabel("Differenz",**font_kwargs)

    # Generate the legend for all subplots
    plt.legend(handles=[ax[0].lines[0], ax[0].lines[1], ax[1].lines[0], ax[1].lines[1], ax[2].lines[0], ax[2].lines[1], ax[2].lines[2], ax[3].lines[0]], loc='upper right')

    plt.tight_layout() # Make sure the legend doesn't cover any plot

    plt.savefig(f"Hahn_gamma={gamma};dt={dt}.png")

for dt in [1e-2, 1e-3, 1e-4]:
    run_sim(0, dt)

for gamma in [0.1, 1, 1.2]:
    run_sim(gamma, 1e-4)