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


# ==============================
# Parameters
# ==============================
nmon = 6.0
nyear = 2.0
dt = 0.5

w = 2.0 * np.pi / nmon
a = 1.0
c = 0.0


# ==============================
# Time coordinate setting
# ==============================
t = np.arange(0.0, nmon*nyear, dt)
nt = len(t)

# ==============================
# Prescribed forcing
# dT/dt = F(t)
# F(t) = a * w * sin(w t)
# ==============================
F = a * w * np.sin(w * t)


# ==============================
# Analytical solution
# T(t) = -a cos(w t) + c
# ==============================
T_ana = -a * np.cos(w * t) + c


# ==============================
# Numerical solutions
# ==============================
T_num1 = np.zeros(nt)

# Initial Conditions
T_num1[0] = 0.0

# Do-loop from n=0 to N-2, obtaining next time step T 
for n in range(nt - 1):
    T_num1[n + 1] = 1.0


# ==============================
# Errors
# ==============================
err1 = T_num1 - T_ana


# ==============================
# Three-panel plot
# ==============================
fig, axes = plt.subplots(3, 1, figsize=(10, 10), sharex=True)

# Panel 1: forcing
axes[0].plot(t, F, 'o', color='k', label='F(t)')
axes[0].set_ylabel('Forcing')
axes[0].set_title('Prescribed Forcing')
axes[0].grid(True)
axes[0].legend()

# Panel 2: analytical and numerical solutions
axes[1].plot(t, T_ana, linewidth=2, color='k', label='Analytical')
axes[1].plot(t, T_num1, 'o', color='C1', label='Numerical')
axes[1].set_ylabel('Temperature')
axes[1].set_title('Analytical and Numerical Solutions')
axes[1].grid(True)
axes[1].legend()

# Panel 3: errors
axes[2].plot(t, err1, 'o', color='C1', label='Numerical error')
axes[2].set_xlabel('Time')
axes[2].set_ylabel('Error')
axes[2].set_title('Numerical Error')
axes[2].set_xticks(np.arange(0, t[-1] + 2.*dt, nmon))
axes[2].grid(True)
axes[2].legend()

plt.tight_layout()


# ----------------------------
# Save figure using script name
# ----------------------------
script_name = os.path.splitext(os.path.basename(__file__))[0]
outfile = script_name + '.png'

plt.savefig(outfile, dpi=150)
print('Figure saved as:', outfile)

plt.show()