import sys
import numpy as np
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt
import matplotlib.animation as animation

# ==========================
# parameters
# ==========================
tmax = 20.0
tstep = 0.05

save = 0
if len(sys.argv) > 1:
    save = int(sys.argv[1])

# solve_ivp の method を CLI で選択
method = "RK45"
if len(sys.argv) > 2:
    method = sys.argv[2]

print(f"Using solve_ivp method = {method}")

# ==========================
# physical constants
# ==========================
g = 9.81
L1 = 1.0
L2 = 1.0
m1 = 1.0
m2 = 1.0

# ==========================
# initial conditions
# ==========================
theta1 = np.pi / 2
theta2 = np.pi / 2
theta1_dot = 0.0
theta2_dot = 0.0

state0 = [theta1, theta2, theta1_dot, theta2_dot]

# ==========================
# ODE system (solve_ivp uses f(t, y))
# ==========================
def derivs(t, state):
    th1, th2, w1, w2 = state
    delta = th2 - th1

    den1 = (m1 + m2) * L1 - m2 * L1 * np.cos(delta)**2
    den2 = (L2 / L1) * den1

    dth1 = w1
    dth2 = w2

    dw1 = (
        m2 * L1 * w1**2 * np.sin(delta) * np.cos(delta)
        + m2 * g * np.sin(th2) * np.cos(delta)
        + m2 * L2 * w2**2 * np.sin(delta)
        - (m1 + m2) * g * np.sin(th1)
    ) / den1

    dw2 = (
        -m2 * L2 * w2**2 * np.sin(delta) * np.cos(delta)
        + (m1 + m2) * g * np.sin(th1) * np.cos(delta)
        - (m1 + m2) * L1 * w1**2 * np.sin(delta)
        - (m1 + m2) * g * np.sin(th2)
    ) / den2

    return [dth1, dth2, dw1, dw2]

# ==========================
# time grid
# ==========================
t_eval = np.arange(0, tmax, tstep)

# ==========================
# solve ODE
# ==========================
sol = solve_ivp(
    derivs,
    t_span=(0, tmax),
    y0=state0,
    t_eval=t_eval,
    method=method
)

y = sol.y.T

# ==========================
# convert to Cartesian coordinates
# ==========================
x1 = L1 * np.sin(y[:, 0])
y1 = -L1 * np.cos(y[:, 0])
x2 = x1 + L2 * np.sin(y[:, 1])
y2 = y1 - L2 * np.cos(y[:, 1])

# ==========================
# animation
# ==========================
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 * tstep))
    return line, time_text

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

if save:
    ani.save('double_pendulum_solveivp.gif', writer='pillow', fps=20)

plt.show()
