import matplotlib.pyplot as plt
import numpy as np
import matplotlib.animation as animation

# Euler
def euler(theta, omega, dt):
    for i in range(1, len(theta)):
        theta[i] = theta[i-1] + omega[i-1] * dt
        omega[i] = omega[i-1] - np.sin(theta[i-1]) * dt
    return theta, omega
        
# Euler-Cromer (semi-implicito)
def euler_cromer(theta, omega, dt):
    for i in range(1, len(theta)):
        omega[i] = omega[i-1] - np.sin(theta[i-1]) * dt
        theta[i] = theta[i-1] + omega[i] * dt
    return theta, omega

# Euler-Richardson
def euler_richardson(theta, omega, dt):
    for i in range(1, len(theta)):
        theta_mid = theta[i-1] + omega[i-1] * 0.5 * dt
        omega_mid = omega[i-1] - np.sin(theta[i-1]) * 0.5 * dt
        theta[i] = theta[i-1] + omega_mid * dt
        omega[i] = omega[i-1] - np.sin(theta_mid) * dt
    return theta, omega

# Euler-Richardson com amortecimento
def euler_richardson_dumping(theta, omega, dt, damp):
    for i in range(1, len(theta)):
        alpha = -np.sin(theta[i-1]) - damp * omega[i-1]
        theta_mid = theta[i-1] + omega[i-1] * 0.5 * dt
        omega_mid = omega[i-1] + alpha  * 0.5 * dt
        alpha_mid = -np.sin(theta_mid) - damp * omega_mid
        theta[i] = theta[i-1] + omega_mid * dt
        omega[i] = omega[i-1] + alpha_mid * dt
    return theta, omega

# Animation
def render(theta, dt):
    # Mass position
    pos = np.zeros((len(theta),2))
    pos[:,0], pos[:,1] = np.sin(theta), -np.cos(theta)

    # Animation
    fig = plt.figure(figsize=(8, 6))
    ax = fig.add_subplot(autoscale_on=False, xlim=(-1.25, 1.25), ylim=(-1.25, 0))
    ax.set_aspect('equal')
    ax.grid()

    line, = ax.plot([], [], 'o-', lw=2)

    def animate(i):
        line.set_data([0, pos[i,0]], [0, pos[i,1]])
        return line, 

    ani = animation.FuncAnimation(
        fig, animate, len(pos), interval=dt*1000, blit=True)

    plt.show()

def main():

    t_stop = 50.0
    dt = 0.01
    t = np.arange(0, t_stop, dt)

    theta = np.zeros(len(t))
    theta[0] = np.pi / 3    # theta0
    omega = np.zeros(len(t))

    #theta, omega = euler(theta, omega, dt)
    #theta, omega = euler_cromer(theta, omega, dt)
    theta, omega = euler_richardson(theta, omega, dt)

    #theta, omega = rk2_dumping(theta, omega, dt, 0.1)

    # Energia total
    max_energy = 1 - np.cos(theta[0])
    total_energy = 1 - np.cos(theta) + 0.5 * omega**2
    energy_error = np.sqrt(np.sum((total_energy-max_energy)**2) / len(theta))
    print('Erro na energia total: ', energy_error)
    print('Energia inicial: ', total_energy[0])
    print('Energia final: ', total_energy[-1])

    # Animation
    render(theta, dt)

    # Plot t x theta
    plt.plot(t, theta)
    plt.show()
    

if __name__ == '__main__': main()