import os
import numpy as np
import xarray as xr
#import matplotlib
#matplotlib.use("Agg")
import matplotlib.pyplot as plt


# ============================================================
# Figure settings
# ============================================================
script_name = os.path.splitext(os.path.basename(__file__))[0]

outfile_png = script_name.replace("_template", "") + ".png"

# ============================================================
# Settings
# ============================================================
FDIR_CTL = "../CTL_Lorenz/out/"
FDIR_PRD = "../PRD_Lorenz/out/"

# Prediction period
TSTR = 50
TEND = 70


# ============================================================
# Read control simulation
# ============================================================
fin_ctl = xr.open_mfdataset(
    FDIR_CTL + "out.*.nc",
    combine="by_coords"
)

time = fin_ctl.time.sel(time=slice(TSTR, TEND))

Xctl = fin_ctl.X.sel(time=slice(TSTR, TEND))
Yctl = fin_ctl.Y.sel(time=slice(TSTR, TEND))
Zctl = fin_ctl.Z.sel(time=slice(TSTR, TEND))


# ============================================================
# Read perturbed simulation
# ============================================================
fin_prd = xr.open_mfdataset(
    FDIR_PRD + "out.*.nc",
    combine="by_coords"
)

Xprd = fin_prd.X.sel(time=slice(TSTR, TEND))
Yprd = fin_prd.Y.sel(time=slice(TSTR, TEND))
Zprd = fin_prd.Z.sel(time=slice(TSTR, TEND))

# Align CTL and PRD using common time coordinates
Xctl, Xprd = xr.align(Xctl, Xprd, join="inner")
Yctl, Yprd = xr.align(Yctl, Yprd, join="inner")
Zctl, Zprd = xr.align(Zctl, Zprd, join="inner")

time = Xctl.time

# ============================================================
# Calculate absolute error
# ============================================================
errX = np.abs(Xprd - Xctl)
errY = np.abs(Yprd - Yctl)
errZ = np.abs(Zprd - Zctl)


# ============================================================
# Print maximum absolute error
# ============================================================
print("Maximum absolute error")
print(f"X : {float(errX.max()):8.4f}")
print(f"Y : {float(errY.max()):8.4f}")
print(f"Z : {float(errZ.max()):8.4f}")


# ============================================================
# Plot error growth
# ============================================================
fig, ax = plt.subplots(3, 1, figsize=(9, 7), sharex=True)

ax[0].plot(time, errX, color="C0", linewidth=1.5)
ax[0].set_ylabel("|X$_{PRD}$ - X$_{CTL}$|")
ax[0].set_title("Initial-Condition Error Growth")

ax[1].plot(time, errY, color="C1", linewidth=1.5)
ax[1].set_ylabel("|Y$_{PRD}$ - Y$_{CTL}$|")

ax[2].plot(time, errZ, color="C2", linewidth=1.5)
ax[2].set_ylabel("|Z$_{PRD}$ - Z$_{CTL}$|")
ax[2].set_xlabel("Time")

for a in ax:
    a.grid(True)

plt.tight_layout()
plt.savefig(outfile_png, dpi=150)
plt.show()
plt.close()

print("Figure saved as:", outfile_png)

fin_ctl.close()
fin_prd.close()