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

# 定数の設定
g = 9.81  # 重力加速度 (m/s^2)
L1 = 1.0  # 上部の振り子の長さ (m)
L2 = 1.0  # 下部の振り子の長さ (m)
m1 = 1.0  # 上部の振り子の質量 (kg)
m2 = 1.0  # 下部の振り子の質量 (kg)

# 初期条件
theta1 = np.pi / 2  # 上部の振り子の初期角度 (ラジアン)
theta2 = np.pi / 2  # 下部の振り子の初期角度 (ラジアン)
theta1_dot = 0.0    # 上部の振り子の初期角速度 (rad/s)
theta2_dot = 0.0    # 下部の振り子の初期角速度 (rad/s)

tmax = 120.0
tstep = 0.05
interval = 10

# 運動方程式の設定
def derivs(state, t):
    dtheta1_dt = state[2]
    dtheta2_dt = state[3]
    
    delta = state[1] - state[0]
    
    denominator1 = (m1 + m2) * L1 - m2 * L1 * np.cos(delta) * np.cos(delta)
    denominator2 = (L2 / L1) * denominator1
    
    dtheta1_dot_dt = ((m2 * L1 * state[2] ** 2 * np.sin(delta) * np.cos(delta) +
                      m2 * g * np.sin(state[1]) * np.cos(delta) +
                      m2 * L2 * state[3] ** 2 * np.sin(delta) -
                      (m1 + m2) * g * np.sin(state[0])) / denominator1)
    
    dtheta2_dot_dt = ((-m2 * L2 * state[3] ** 2 * np.sin(delta) * np.cos(delta) +
                      (m1 + m2) * g * np.sin(state[0]) * np.cos(delta) -
                      (m1 + m2) * L1 * state[2] ** 2 * np.sin(delta) -
                      (m1 + m2) * g * np.sin(state[1])) / denominator2)
    
    return [dtheta1_dt, dtheta2_dt, dtheta1_dot_dt, dtheta2_dot_dt]

# 初期状態と時間配列の設定
state = [theta1, theta2, theta1_dot, theta2_dot]
t = np.arange(0.0, tmax, tstep)

# 運動方程式を数値的に解く
from scipy.integrate import odeint
y = odeint(derivs, state, t)

# シミュレーション結果の表示
x1 = L1 * np.sin(y[:, 0])
y1 = -L1 * np.cos(y[:, 0])
x2 = L2 * np.sin(y[:, 1]) + x1
y2 = -L2 * np.cos(y[:, 1]) + y1

fig = plt.figure()
ax = fig.add_subplot(111, autoscale_on=False, xlim=(-2, 2), ylim=(-2, 2))
ax.set_aspect('equal')
ax.grid()

line, = ax.plot([], [], 'o-', lw=2)
time_template = 'time = %.1fs'
time_text = ax.text(0.05, 0.9, '', transform=ax.transAxes)

def init():
    line.set_data([], [])
    time_text.set_text('')
    return line, time_text

def animate(i):
    thisx = [0, x1[i], x2[i]]
    thisy = [0, y1[i], y2[i]]
    line.set_data(thisx, thisy)
    time_text.set_text(time_template % (i * 0.05))
    return line, time_text

ani = animation.FuncAnimation(fig, animate, np.arange(1, len(y)),
                              interval=interval, blit=True, init_func=init)

plt.show()