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

# ===== 設定 =====
#変更箇所はここだけ
T_sync =43.4   #何秒後に揃うかの時間(s)
num = 15       #振り子の数(個)
n_start = 25   #一番長い振り子の振動回数
#=================

g = 980.0      #重力加速度(cm/s^2)

n_values = np.arange(n_start, n_start + num)

# 周期と長さ
periods = T_sync / n_values
lengths = g * (periods / (2 * np.pi))**2

# 角振動数
omega = 2 * np.pi / periods

# 表示
print(f"{T_sync}秒後に再びそろう振り子の長さ")
for i, (n, T, L) in enumerate(zip(n_values, periods, lengths), start=1):
    print(f"{i:2d}番目: {n:2d}回振動, 周期={T:.4f} s, 長さ={L:.2f} cm")

# ===== アニメーション設定 =====
fig, ax = plt.subplots()

ax.set_xlim(-40, 40)
ax.set_ylim(0, max(lengths) * 1.2)
ax.invert_yaxis()
ax.set_xlabel("x")
ax.set_ylabel("length [cm]")
ax.set_title("Pendulum Wave")

balls, = ax.plot([], [], "o", markersize=8)
lines = [ax.plot([], [], "-")[0] for _ in range(num)]
time_text = ax.text(0.02, 0.95, "", transform=ax.transAxes)

dt = 0.02
times = np.arange(0, T_sync + dt, dt)

# 振幅
amplitude = 20.0


def update(frame):
    t = times[frame]

    # 各振り子の横方向変位
    x = amplitude * np.cos(omega * t)
    y = lengths

    balls.set_data(x, y)

    for i in range(num):
        lines[i].set_data([0, x[i]], [0, y[i]])

    time_text.set_text(f"t = {t:.2f} s")

    return [balls, time_text] + lines


ani = FuncAnimation(
    fig,
    update,
    frames=len(times),
    interval=20,
    blit=True
)

plt.show()