# ============================================================
# ACC of X at different forecast lead times
#
# Control simulation:
#   ../CTL/out/
#
# Hindcast experiments:
#   ../HINDCAST/0055/out/
#   ../HINDCAST/0060/out/
#   ../HINDCAST/0065/out/
#   ../HINDCAST/0070/out/
#   ../HINDCAST/0075/out/
#   ../HINDCAST/0080/out/
#   ../HINDCAST/0085/out/
#   ../HINDCAST/0090/out/
#
# Five panels show lead times:
#   0, 1, 2, 3, and 4
# ============================================================

import os
import numpy as np
import xarray as xr
import matplotlib.pyplot as plt


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

script_name = os.path.splitext(os.path.basename(__file__))[0]
outfile_png = script_name + ".png"

figure_title = "Lorenz Model: ACC of X"

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

# Axis limits
YLIM_X = (-25.0, 25.0)

# Forecast lead times
LEAD_TIMES = [0, 1, 2, 3, 4]


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

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


# ============================================================
# Hindcast initialization times
# ============================================================

HINDCAST_EXPERIMENTS = [
    "0055",
    "0060",
    "0065",
    "0070",
    "0075",
    "0080",
    "0085",
    "0090",
]


# ============================================================
# 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)
)


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

fig, ax = plt.subplots(
    5,
    1,
    figsize=(10, 12),
    sharex=True
)

fig.suptitle(
    figure_title,
    fontsize=14
)


# ============================================================
# Plot the control simulation in all five panels
# ============================================================

for ilead in range(len(LEAD_TIMES)):

    ax[ilead].plot(
        time_ctl,
        Xctl,
        color="black",
        linewidth=1.5,
        label="Control" if ilead == 0 else None
    )

    ax[ilead].set_ylabel("X")
    ax[ilead].set_ylim(YLIM_X)
    ax[ilead].set_xlim(TSTR_CTL, TEND_CTL)
    ax[ilead].grid(True)

    ax[ilead].annotate(
        f"Lead time = {LEAD_TIMES[ilead]}",
        xy=(0.02, 0.82),
        xycoords="axes fraction",
        backgroundcolor="white",
        fontsize=9
    )

ax[-1].set_xlabel("Model time")


# ============================================================
# Arrays for forecast verification
# ============================================================

NLEAD = len(LEAD_TIMES)
NEXP = len(HINDCAST_EXPERIMENTS)

Xtrue = np.zeros((NLEAD, NEXP))
Xforecast = np.zeros((NLEAD, NEXP))


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

for iexp in range(NEXP):

    experiment = HINDCAST_EXPERIMENTS[iexp]

    initial_time = float(int(experiment))

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

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

    time_prd = fin_prd.time
    Xprd = fin_prd.X

    # --------------------------------------------------------
    # Repeat the spaghetti plot in all panels
    # --------------------------------------------------------

    for ilead in range(NLEAD):

        lead_time = LEAD_TIMES[ilead]

        verification_time = initial_time + lead_time

        # Add a legend label only once
        if iexp == 0 and ilead == 0:
            hindcast_label = "Hindcasts"
        else:
            hindcast_label = None

        ax[ilead].plot(
            time_prd,
            Xprd,
            color="C1",
            linewidth=1.0,
            alpha=0.8,
            label=hindcast_label
        )

        # ----------------------------------------------------
        # Extract true and forecast states
        # ----------------------------------------------------

        true_value = float(
            fin_ctl.X.sel(time=verification_time)
        )

        forecast_value = float(
            Xprd.sel(time=verification_time)
        )

        Xtrue[ilead, iexp] = true_value
        Xforecast[ilead, iexp] = forecast_value

        # ----------------------------------------------------
        # Connect the true and forecast states
        # ----------------------------------------------------

        ax[ilead].plot(
            [verification_time, verification_time],
            [true_value, forecast_value],
            color="0.5",
            linewidth=0.8
        )

        # ----------------------------------------------------
        # Mark the true state
        # ----------------------------------------------------

        if iexp == 0 and ilead == 0:
            true_label = "True state"
        else:
            true_label = None

        ax[ilead].plot(
            verification_time,
            true_value,
            marker="o",
            markersize=6,
            color="C0",
            markeredgecolor="black",
            linestyle="None",
            label=true_label
        )

        # ----------------------------------------------------
        # Mark the forecast state
        # ----------------------------------------------------

        if iexp == 0 and ilead == 0:
            forecast_label = "Forecast state"
        else:
            forecast_label = None

        ax[ilead].plot(
            verification_time,
            forecast_value,
            marker="x",
            markersize=7,
            color="C3",
            markeredgewidth=2.0,
            linestyle="None",
            label=forecast_label
        )

    fin_prd.close()


# ============================================================
# Calculate ACC
# ============================================================

ACC = np.zeros(NLEAD)

print("========================================")
print("Anomaly Correlation Coefficient")
print("========================================")

for ilead in range(NLEAD):

    ACC[ilead] = np.corrcoef(
        Xtrue[ilead, :],
        Xforecast[ilead, :]
    )[0, 1]

    print(
        f"Lead time = {LEAD_TIMES[ilead]}, "
        f"ACC = {ACC[ilead]:.3f}"
    )

    ax[ilead].annotate(
        f"ACC = {ACC[ilead]:.3f}",
        xy=(0.98, 0.82),
        xycoords="axes fraction",
        horizontalalignment="right",
        backgroundcolor="white",
        fontsize=10
    )


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

ax[0].legend(
    loc="lower right",
    fontsize=8,
    ncol=2
)

plt.tight_layout(
    rect=[0.0, 0.0, 1.0, 0.97]
)


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

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

plt.show()
plt.close()

print("Figure saved as:", outfile_png)

fin_ctl.close()