# Blatt 9 Aufgabe 5

import numpy as np
import matplotlib.pyplot as plt

# Const
m1 = 1.0  # kg
m2 = 1.0  # kg
l = 1.0   # m
g = 9.81  # m/s^2

# Params
T = 30.0#s
dt = 1e-4  # Timestep
dt2 = 1e-5  # Timestep 2 (very long times)

# Initial conditions
theta1_0 = np.pi / 2
theta2_0 = np.pi
omega1_0 = 0.0
omega2_0 = 0.0

def a(theta, omega):
    # Unpack the params
    theta1, theta2 = theta
    omega1, omega2 = omega
    delta = theta1 - theta2

    # Represent the equation in terms of M and Q
    # L, m1, m2 not used for the matrix equation ??
    # Do I need to use the result from task a) ?
    M = np.array([
        [2, np.cos(delta)],
        [np.cos(delta), 1]
    ])
    Q = np.array([
        -omega2**2 * np.sin(delta) - 2 * g * np.sin(theta1),
        omega1**2 * np.sin(delta) - g * np.sin(theta2)
    ])

    # TODO: better method?
    return np.linalg.solve(M, Q) # Return the acc in a theta double dot vector

# Run the main simulation
def run_simulation(theta0, omega0, dt):
    # Setup state arrays
    times = np.arange(0, T, dt)
    theta = np.zeros((len(times), 2))
    omega = np.zeros((len(times), 2))
    energy = np.zeros(len(times))

    # Initialize
    theta[0] = theta0
    omega[0] = omega0
    acc = a(theta[0], omega[0])
    print(acc)

    # Iterate
    for i in range(1, len(times)):
        # VV-step
        theta[i] = theta[i-1] + omega[i-1] * dt + 0.5 * acc * dt**2
        a_new = a(theta[i], omega[i-1])
        omega[i] = omega[i-1] + 0.5 * (acc + a_new) * dt
        acc = a_new

        # Energy
        theta1, theta2 = theta[i]
        omega1, omega2 = omega[i]
        E_kin = 0.5 * m1 * (l * omega1)**2 + 0.5 * m2 * (
            (l * omega1)**2 + (l * omega2)**2 +
            2 * l**2 * omega1 * omega2 * np.cos(theta1 - theta2)
        )
        E_pot = - (m1 + m2) * g * l * np.cos(theta1) - m2 * g * l * np.cos(theta2)

        # Save the total energy
        energy[i] = E_kin + E_pot

    # Return timeline
    return times, theta, omega, energy

# Start simulation
theta0 = [theta1_0, theta2_0]
omega0 = [omega1_0, omega2_0]
times, theta, _, energy = run_simulation(theta0, omega0, dt)

theta0 = [theta1_0 + 10 ** -7, theta2_0]
omega0 = [omega1_0, omega2_0]
_, theta2, _, _ = run_simulation(theta0, omega0, dt)

theta0 = [theta1_0, theta2_0]
omega0 = [omega1_0, omega2_0]
times3, theta3, _, energy3 = run_simulation(theta0, omega0, dt2) # Smaller timestep

# Quick plotting

# Energy
plt.figure(figsize=(12, 6))
plt.plot(times, energy, label='Total energy', color='blue')
plt.plot(times3, energy3, label='Total energy smaller dt', color='blue', linestyle="--")
plt.xlabel('Time [s]')
plt.ylabel('Energy [J]')
plt.title('Total energy of the double pendulum')
plt.grid(True)
plt.legend()
plt.tight_layout()
plt.savefig("double_energy.svg")

# Angles
plt.figure(figsize=(12, 6))

# Reduce angles
theta_mod = (theta + np.pi) % (2 * np.pi) - np.pi
theta2_mod = (theta2 + np.pi) % (2 * np.pi) - np.pi

# Plot with module (less information)
#plt.figure(figsize=(12, 6))
#plt.plot(times, theta_mod[:, 0], label=r'$\theta_1(t)$', color='red')
#plt.plot(times, theta_mod[:, 1], label=r'$\theta_2(t)$', color='green')
#plt.plot(times, theta2_mod[:, 0], label=r'$\theta_1(t)$ (small change)', color='blue')
#plt.plot(times, theta2_mod[:, 1], label=r'$\theta_2(t)$ (small change)', color='gold')

plt.plot(times, theta[:, 0] , label=r'$\theta_1(t)$', color='red')
plt.plot(times, theta[:, 1] , label=r'$\theta_2(t)$', color='green')
plt.plot(times3, theta3[:, 0] , label=r'$\theta_1(t)$ smaller dt', color='red', linestyle="--")
plt.plot(times3, theta3[:, 1] , label=r'$\theta_2(t)$ smaller dt', color='green', linestyle="--")
plt.plot(times, theta2[:, 0] , label=r'$\theta_1(t)$ changed', color='blue')
plt.plot(times, theta2[:, 1] , label=r'$\theta_2(t)$ changed', color='yellow')

plt.xlabel('Time [s]')
plt.ylabel('Angle [rad]')
plt.title('Trajectories of both pendulums')
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.savefig("double_angles.svg")