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"

# Model output directories
FDIR_CTL = "../CTL_Lorenz/out/"
FDIR_PRD = "../PRD_Lorenz/out/"

# Use a longer period for the control simulation
TSTR_CTL = 0
TEND_CTL = 100

# Use the prediction period for the perturbed simulation
TSTR_PRD = 50
TEND_PRD = 70


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

time_ctl = fin_ctl.time.sel(time=slice(TSTR_CTL, TEND_CTL))
Xctl = fin_ctl.X.sel(time=slice(TSTR_CTL, TEND_CTL))
Yctl = fin_ctl.Y.sel(time=slice(TSTR_CTL, TEND_CTL))
Zctl = fin_ctl.Z.sel(time=slice(TSTR_CTL, TEND_CTL))


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

time_prd = fin_prd.time.sel(time=slice(TSTR_PRD, TEND_PRD))
Xprd = fin_prd.X.sel(time=slice(TSTR_PRD, TEND_PRD))
Yprd = fin_prd.Y.sel(time=slice(TSTR_PRD, TEND_PRD))
Zprd = fin_prd.Z.sel(time=slice(TSTR_PRD, TEND_PRD))

# ============================================================
# Plot Lorenz control and perturbed simulations
# ============================================================
fig, ax = plt.subplots(3, 1, figsize=(9, 7), sharex=True)

ax[0].plot( time_ctl, Xctl, color="C0", linewidth=2.0, label="CTL")
ax[0].plot( time_prd, Xprd, color="C1", linewidth=1.0, label="PRD")

ax[0].set_ylabel("X")
ax[0].set_title("Control Simulation and Prediction")
ax[0].legend(loc="best", fontsize=8)

# ----------------------------
# Y
# ----------------------------
ax[1].plot( time_ctl, Yctl, color="C0", linewidth=2.0, label="CTL")
ax[1].plot( time_prd, Yprd, color="C1", linewidth=1.0, label="PRD")
ax[1].set_ylabel("Y")

# ----------------------------
# Z
# ----------------------------
ax[2].plot( time_ctl, Zctl, color="C0", linewidth=2.0, label="CTL")
ax[2].plot( time_prd, Zprd, color="C1", linewidth=1.0, label="PRD")


ax[2].set_ylabel("Z")
ax[2].set_xlabel("Time")

# ----------------------------
# Common plot settings
# ----------------------------
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()