x = np.loadtxt(f"{data_dir}/x_grid.dat") # Load data for animation
eta_files = sorted([f for f in os.listdir(data_dir) if f.startswith("eta_")])
u_files   = sorted([f for f in os.listdir(data_dir) if f.startswith("u_")])
nt = len(eta_files)
eta_array = np.zeros((nt, nx))
u_array   = np.zeros((nt, nx))
for i in range(nt):
    eta_array[i] = np.loadtxt(f"{data_dir}/{eta_files[i]}")
    u_array[i]   = np.loadtxt(f"{data_dir}/{u_files[i]}")
bottom_display = -0.2
fig, axes = plt.subplots(3,1, figsize=(15,8), sharex=True)
fig.suptitle("1D Shallow Water Basin (Wind Forced)", fontsize=12)
ax0, ax1, ax2 = axes
surface_line, = ax1.plot(x/1000, eta_array[0], color="midnightblue", lw=2)
vel_line,     = ax2.plot([], [], lw=2, color="hotpink")
def update(frame):
    eta = eta_array[frame]
    u   = u_array[frame]
    surface_line.set_data(x/1000, eta)
    vel_line.set_data(x/1000, u)
    time_hours = time_array[frame] / 3600
    ax1.set_title(f"Surface Evolution (t = {time_hours:.2f} hours)")
    ax2.set_title(f"Velocity (t = {time_hours:.2f} hours)")
    return surface_line, vel_line
ani = animation.FuncAnimation(
    fig,
    update,
    frames=min(nt,200),
    interval=120
)
ani.save("swe_exaggerated_surface.gif", writer="pillow", fps=8)
print("Saved: swe_exaggerated_surface.gif")
plt.close(fig)