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 + ".png"

FDIR = "../out/"
TSTR = 60
TEND = 240


# ============================================================
# Read model output
# ============================================================
fin = xr.open_mfdataset(
    FDIR + "out.*.nc",
    combine="by_coords"
)

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

X = fin.X.sel(time=slice(TSTR, TEND))
Y = fin.Y.sel(time=slice(TSTR, TEND))
Z = fin.Z.sel(time=slice(TSTR, TEND))


# ============================================================
# Plot total climate state
# ============================================================
fig, ax = plt.subplots(3, 1, figsize=(9, 7), sharex=True)

ax[0].plot(time, X, color="C0", linewidth=1.0)
ax[0].set_ylabel("X")
ax[0].set_title("Lorenz Climate Model")

ax[1].plot(time, Y, color="C0", linewidth=1.0)
ax[1].set_ylabel("Y")

ax[2].plot(time, Z, color="C0", linewidth=1.0)
ax[2].set_ylabel("Z")
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.close()