#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Dec 13 15:48:00 2024

@author: laurentdumas
"""
import numpy as np
import matplotlib.pyplot as plt
k = -0.001             
t_max = 300        
dt = 0.01          
theta1_0 = np.pi / 6 
theta2_0 = 0.0        
v1_0 = 0.0            
v2_0 = 0.0            

def pendules_couples(t, Y, k):
    theta1, v1, theta2, v2 = Y  
    dtheta1 = v1
    dv1 = -np.sin(theta1) + k * (theta2 - theta1)
    dtheta2 = v2
    dv2 = -np.sin(theta2) + k * (theta1 - theta2)
    return np.array([dtheta1, dv1, dtheta2, dv2])

def runge_kutta4(f, t0, Y0, t_max, dt, k):
    t = np.arange(t0, t_max, dt)
    Y = np.zeros((len(t), len(Y0)))
    Y[0] = Y0
    
    for i in range(1, len(t)):
        k1 = dt * f(t[i-1], Y[i-1], k)
        k2 = dt * f(t[i-1] + dt/2, Y[i-1] + k1/2, k)
        k3 = dt * f(t[i-1] + dt/2, Y[i-1] + k2/2, k)
        k4 = dt * f(t[i-1] + dt, Y[i-1] + k3, k)
        Y[i] = Y[i-1] + (k1 + 2*k2 + 2*k3 + k4) / 6
        
    return t, Y

Y0 = [theta1_0, v1_0, theta2_0, v2_0]

t, Y = runge_kutta4(pendules_couples, 0, Y0, t_max, dt, k)
plt.figure(1)
plt.plot(t, Y[:, 0], label=r"$\theta_1(t)$", color="b")
plt.figure(2)
plt.plot(t, Y[:, 2], label=r"$\theta_2(t)$", color="r")
plt.show()