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


# ==============================
# Parameters
# ==============================
ntime = 50
dt = 1.0

# Initial condition
T0 = 5.0
t0 = 0.0

# ==============================
# Forcing term
# ==============================
def force_term(T):
    F = 1.0
    return F


# ==============================
# Time integration
# ==============================
t_save = np.zeros(ntime)
T_save = np.zeros(ntime)
F_save = np.zeros(ntime)

# Initial values
t = t0
T = T0

t_save[0] = t
T_save[0] = T
F_save[0] = force_term(T)

# Do-loop
for it in range(ntime - 1):
    F = force_term(T)
    t = t + dt
    T = T + dt * F

    t_save[it + 1] = t
    T_save[it + 1] = T
    F_save[it + 1] = F

# ==============================
# Two-panel figure
# ==============================
fig, axes = plt.subplots(2, 1, figsize=(10, 7), sharex=True)

lineF, = axes[0].plot([], [], 'o-', color='k', label='Forcing')
pointF, = axes[0].plot([], [], 'ro')
axes[0].set_ylabel('F')
axes[0].set_title('Forcing')
axes[0].grid(True)
axes[0].legend()

lineT, = axes[1].plot([], [], 'o-', color='C1', label='Temperature')
pointT, = axes[1].plot([], [], 'ro')
axes[1].set_xlabel('Time')
axes[1].set_ylabel('T')
axes[1].set_title('Temperature')
axes[1].grid(True)
axes[1].legend()

# Axis ranges
fmin = np.min(F_save)
fmax = np.max(F_save)
tmin = np.min(T_save)
tmax = np.max(T_save)

if fmax == fmin:
    fpad = 0.5
else:
    fpad = 0.1 * (fmax - fmin)

if tmax == tmin:
    tpad = 0.5
else:
    tpad = 0.1 * (tmax - tmin)

axes[0].set_xlim(t_save[0], t_save[-1])
axes[0].set_ylim(fmin - fpad, fmax + fpad)

axes[1].set_xlim(t_save[0], t_save[-1])
axes[1].set_ylim(tmin - tpad, tmax + tpad)


# ==============================
# Animation
# ==============================
def init():
    lineF.set_data([], [])
    pointF.set_data([], [])
    lineT.set_data([], [])
    pointT.set_data([], [])
    return lineF, pointF, lineT, pointT


def update(frame):
    lineF.set_data(t_save[:frame + 1], F_save[:frame + 1])
    pointF.set_data([t_save[frame]], [F_save[frame]])

    lineT.set_data(t_save[:frame + 1], T_save[:frame + 1])
    pointT.set_data([t_save[frame]], [T_save[frame]])

    return lineF, pointF, lineT, pointT


ani = animation.FuncAnimation(
    fig,
    update,
    frames=ntime,
    init_func=init,
    interval=300,
    blit=True
)

plt.tight_layout()


# ----------------------------
# Save animation using script name
# ----------------------------
script_name = os.path.splitext(os.path.basename(__file__))[0]
outfile_mp4 = script_name + '.mp4'

writer = animation.FFMpegWriter(fps=4)
ani.save(outfile_mp4, writer=writer)
print('Animation saved as:', outfile_mp4)


# ----------------------------
# Save final figure
# ----------------------------
outfile_png = script_name + '.png'
plt.savefig(outfile_png, dpi=150)
print('Figure saved as:', outfile_png)

plt.show()