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 = 0
TEND = 120


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

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


# ============================================================
# Check debug variables
# ============================================================
debug_vars = ["Xint", "Yint", "Zint", "Xext", "Yext", "Zext"]

for v in debug_vars:
    if v not in fin:
        raise ValueError(
            f"{v} is not found in the output file. "
            "Please rerun the model with DEBUG_OUTPUT=True."
        )


# ============================================================
# Plot total, internal, and external components
# ============================================================
varlist = ["X", "Y", "Z"]

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

for i, var in enumerate(varlist):

    total = fin[var].sel(time=slice(TSTR, TEND))
    internal = fin[var + "int"].sel(time=slice(TSTR, TEND))
    external = fin[var + "ext"].sel(time=slice(TSTR, TEND))

    ax[i].plot(time, total, color="C0", linewidth=1.2, label=f"{var} total")
    ax[i].plot(time, internal, color="C1", linewidth=1.0, label=f"{var} internal")
    ax[i].plot(time, external, color="C2", linewidth=1.0, label=f"{var} external")

    ax[i].set_ylabel(var)
    ax[i].grid(True)
    ax[i].legend(loc="upper right", fontsize=8)

ax[0].set_title("Lorenz Climate Model: Total, Internal, and External Components")
ax[-1].set_xlabel("Time")

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

print("Figure saved as:", outfile_png)

fin.close()