# ============================================================
# Spaghetti plot of control and hindcast experiments
#
# Control simulation:
#   ../CTL/out/
#
# Hindcast experiments:
#   ../HINDCAST/0055/out/
#   ../HINDCAST/0060/out/
#   ../HINDCAST/0065/out/
# ============================================================

import os
import glob

import xarray as xr
import matplotlib.pyplot as plt


# ============================================================
# Figure settings
# ============================================================

# Use the Python script name as the output figure name
script_name = os.path.splitext(os.path.basename(__file__))[0]
outfile_png = script_name + ".png"

figure_title = "Lorenz Model: Control and Hindcast Experiments"

# Control simulation period shown in the figure
TSTR_CTL = 50
TEND_CTL = 75

# Axis limits
YLIM_X = (-25.0, 25.0)
YLIM_Y = (-25.0, 25.0)
YLIM_Z = (0.0, 60.0)


# ============================================================
# Model output directories
# ============================================================

FDIR_CTL = "../CTL/out/"
FDIR_HINDCAST = "../HINDCAST/dddd/out/"

# Hindcast initialization times
HINDCAST_EXPERIMENTS = [
    "0055",
    "0060",
    "0065",
]


# ============================================================
# Function to read model output
# ============================================================

def read_model_output(directory):
    """
    Read all NetCDF model output files from one directory.

    Parameters
    ----------
    directory : str
        Directory containing out.*.nc files.

    Returns
    -------
    xarray.Dataset
        Combined model output dataset.
    """

    file_pattern = os.path.join(directory, "out.*.nc")
    file_list = sorted(glob.glob(file_pattern))

    if not file_list:
        raise FileNotFoundError(
            f"No model output files were found:\n{file_pattern}"
        )

    print(f"Reading {len(file_list)} files from:")
    print(f"  {directory}")

    dataset = xr.open_mfdataset(
        file_list,
        combine="by_coords",
    )

    # Load the data into memory before closing the NetCDF files
    dataset.load()

    return dataset


# ============================================================
# Read the control simulation
# ============================================================

print("========================================")
print("Read control simulation")
print("========================================")

ds_ctl = read_model_output(FDIR_CTL)

ds_ctl = ds_ctl.sel(time=slice(TSTR_CTL, TEND_CTL))

time_ctl = ds_ctl["time"]
Xctl = ds_ctl["X"]
Yctl = ds_ctl["Y"]
Zctl = ds_ctl["Z"]


# ============================================================
# Create the figure
# ============================================================

fig, axes = plt.subplots(
    nrows=3,
    ncols=1,
    figsize=(10, 8),
    sharex=True,
)

# ------------------------------------------------------------
# Plot the control simulation
# ------------------------------------------------------------

axes[0].plot(
    time_ctl,
    Xctl,
    color="black",
    linewidth=2.0,
    label="Control",
)

axes[1].plot(
    time_ctl,
    Yctl,
    color="black",
    linewidth=2.0,
)

axes[2].plot(
    time_ctl,
    Zctl,
    color="black",
    linewidth=2.0,
)


# ============================================================
# Read and plot the hindcast experiments
# ============================================================

print("========================================")
print("Read hindcast experiments")
print("========================================")

for experiment in HINDCAST_EXPERIMENTS:

    hindcast_directory = FDIR_HINDCAST.replace("dddd", experiment)

    ds_hindcast = read_model_output(hindcast_directory)

    time_hindcast = ds_hindcast["time"]
    Xhindcast = ds_hindcast["X"]
    Yhindcast = ds_hindcast["Y"]
    Zhindcast = ds_hindcast["Z"]

    # Add a legend label only to the first hindcast
    if experiment == HINDCAST_EXPERIMENTS[0]:
        hindcast_label = "Hindcasts"
    else:
        hindcast_label = None

    axes[0].plot(
        time_hindcast,
        Xhindcast,
        color="C1",
        linewidth=1.5,
        label=hindcast_label,
    )

    axes[1].plot(
        time_hindcast,
        Yhindcast,
        color="C1",
        linewidth=1.5,
    )

    axes[2].plot(
        time_hindcast,
        Zhindcast,
        color="C1",
        linewidth=1.5,
    )

    # --------------------------------------------------------
    # Mark the forecast initialization time
    # --------------------------------------------------------

    axes[0].plot(
        time_hindcast[0],
        Xhindcast[0],
        marker="o",
        markersize=6,
        color="C1",
        markeredgecolor="black",
    )

    axes[1].plot(
        time_hindcast[0],
        Yhindcast[0],
        marker="o",
        markersize=6,
        color="C1",
        markeredgecolor="black",
    )

    axes[2].plot(
        time_hindcast[0],
        Zhindcast[0],
        marker="o",
        markersize=6,
        color="C1",
        markeredgecolor="black",
    )

    ds_hindcast.close()


# ============================================================
# Figure formatting
# ============================================================

axes[0].set_title(figure_title)

axes[0].set_ylabel("X")
axes[1].set_ylabel("Y")
axes[2].set_ylabel("Z")
axes[2].set_xlabel("Time")

axes[0].set_xlim(TSTR_CTL, TEND_CTL)

axes[0].set_ylim(YLIM_X)
axes[1].set_ylim(YLIM_Y)
axes[2].set_ylim(YLIM_Z)

for ax in axes:
    ax.grid(True)

axes[0].legend(loc="upper right")

plt.tight_layout()


# ============================================================
# Save the figure
# ============================================================

plt.savefig(
    outfile_png,
    dpi=150,
    bbox_inches="tight",
)

plt.close(fig)

ds_ctl.close()

print("========================================")
print("Figure saved")
print("========================================")
print(outfile_png)