import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.patches import Circle

L1 = 1.0
L2 = 0.8
m1 = 1.0
m2 = 1.5
g = 9.8
k = 0.1

initial_state = np.radians([120.0, 0.0, -20.0, 0.0])

t = np.linspace(0, 20, 1000)

def equations(state, t):
    theta1, z1, theta2, z2 = state
   
    perturbation1 = k * np.sin(5*t) * np.cos(3*theta1)
    perturbation2 = k * np.cos(4*t) * np.sin(2*theta2)
   
    c = np.cos(theta1 - theta2)
    s = np.sin(theta1 - theta2)
   
    denominator = L1 * L2 * (m1 + m2*np.sin(theta1 - theta2)**2)
   
    dz1 = ((-g*(m1 + m2)*L2*np.sin(theta1) + m2*g*L2*np.sin(theta2)*c
           + m2*L2*z2**2*s + m2*L1*z1**2*s*c) / denominator) + perturbation1
   
    dz2 = ((-m2*g*L1*np.sin(theta2) + (m1 + m2)*g*L1*np.sin(theta1)*c
           - (m1 + m2)*L1*z1**2*s - m2*L2*z2**2*s*c) / denominator) + perturbation2
   
    dz1 -= 0.01 * z1
    dz2 -= 0.01 * z2
   
    return [z1, dz1, z2, dz2]

solution = odeint(equations, initial_state, t)
theta1 = solution[:, 0]
theta2 = solution[:, 2]

x1 = L1 * np.sin(theta1)
y1 = -L1 * np.cos(theta1)
x2 = x1 + L2 * np.sin(theta2)
y2 = y1 - L2 * np.cos(theta2)

fig, ax = plt.subplots(figsize=(8, 8))
ax.set_xlim(-(L1+L2+0.2), L1+L2+0.2)
ax.set_ylim(-(L1+L2+0.2), L1+L2+0.2)
ax.set_aspect('equal')
ax.grid()

line, = ax.plot([], [], 'o-', lw=2, color='black')
trace, = ax.plot([], [], '-', lw=1, color='red', alpha=0.5)
ball1 = Circle((0, 0), 0.05, fc='blue', zorder=3)
ball2 = Circle((0, 0), 0.05, fc='green', zorder=3)
ax.add_patch(ball1)
ax.add_patch(ball2)

history_x, history_y = [], []
max_history = 500

def init():
    line.set_data([], [])
    trace.set_data([], [])
    ball1.center = (0, 0)
    ball2.center = (0, 0)
    return line, trace, ball1, ball2

def animate(i):
    points_x = [0, x1[i], x2[i]]
    points_y = [0, y1[i], y2[i]]
   
    line.set_data(points_x, points_y)
    ball1.center = (x1[i], y1[i])
    ball2.center = (x2[i], y2[i])
   
    history_x.append(x2[i])
    history_y.append(y2[i])
    if len(history_x) > max_history:
        history_x.pop(0)
        history_y.pop(0)
    trace.set_data(history_x, history_y)
   
    return line, trace, ball1, ball2

ani = animation.FuncAnimation(
    fig, animate, frames=len(t),
    interval=1000/60, init_func=init, blit=True
)

plt.xlabel('X')
plt.ylabel('Y')
plt.tight_layout()
plt.show()