# ============================================================
# Hindcast skill evaluation
#
# This script:
#   1. Reads the control simulation as the true state.
#   2. Reads all hindcast experiments.
#   3. Converts the data to:
#
#          Xf(t_initial, t_lead)
#          Xtrue(t_initial, t_lead)
#
#      and similarly for Y and Z.
#
#   4. Calculates:
#          ACC
#          RMSE
#          RMSE Skill Score
#
#   5. Creates a 3 x 3 panel figure.
#
# Rows:
#   X, Y, Z
#
# Columns:
#   ACC, RMSE, RMSE Skill Score
# ============================================================

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.replace("_template", "") + ".png"

figure_title = "Lorenz Model Hindcast Skill"


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

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

# ============================================================
# Hindcast experiment settings
# ============================================================

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

NEXP = len(HINDCAST_EXPERIMENTS)

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

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

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

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


FORECAST_LENGTH = 5.0
lead_time_ctl = time_ctl.sel( time=slice( TSTR_CTL, TSTR_CTL + FORECAST_LENGTH)) 
LEAD_TIMES = ( lead_time_ctl.values[:-1] - TSTR_CTL)

NLEAD = len(LEAD_TIMES)

print("Lead-time information")
print("----------------------------")
print("Number of lead times :", NLEAD)
print("First lead time      :", LEAD_TIMES[0])
print("Last lead time       :", LEAD_TIMES[-1])
print("Lead-time interval   :", LEAD_TIMES[1] - LEAD_TIMES[0])


# ============================================================
# Control-run climatology
# ============================================================
#
# The climatological mean is used as the reference forecast
# for calculating the RMSE Skill Score.
# ============================================================

Xclim = float(Xctl.mean())
Yclim = float(Yctl.mean())
Zclim = float(Zctl.mean())

print("Control-run climatology")
print("----------------------------")
print(f"X climatology = {Xclim:.3f}")
print(f"Y climatology = {Yclim:.3f}")
print(f"Z climatology = {Zclim:.3f}")


# ============================================================
# Create forecast and truth arrays
# ============================================================
#
# Array dimensions:
#
#   first dimension  = initialization experiment
#   second dimension = forecast lead time
#
# For example:
#
#   Xf[iexp, ilead]
#   Xtrue[iexp, ilead]
# ============================================================

Xf = np.full((NEXP, NLEAD), np.nan)
Yf = np.full((NEXP, NLEAD), np.nan)
Zf = np.full((NEXP, NLEAD), np.nan)

Xtrue = np.full((NEXP, NLEAD), np.nan)
Ytrue = np.full((NEXP, NLEAD), np.nan)
Ztrue = np.full((NEXP, NLEAD), np.nan)


# ============================================================
# Read hindcast experiments
# ============================================================

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

# Control-output time interval
DT_CTL = float(time_ctl.values[1] - time_ctl.values[0])

for iexp in range(NEXP):

    experiment = HINDCAST_EXPERIMENTS[iexp]

    # Convert "0055" to 55.0
    initial_time = float(int(experiment))

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

    print( f"Experiment {experiment}: " f"T_INITIAL = {initial_time:.0f}")

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

    time_prd = fin_prd.time
    Xprd = fin_prd.X
    Yprd = fin_prd.Y
    Zprd = fin_prd.Z

    # --------------------------------------------------------
    # Confirm that the hindcast contains enough time steps
    # --------------------------------------------------------

    if time_prd.size < NLEAD:
        raise ValueError(
            f"Experiment {experiment} contains only "
            f"{time_prd.size} time steps, but "
            f"{NLEAD} time steps are required."
        )

    # --------------------------------------------------------
    # Store forecast states
    #
    # Xf(t_initial, t_lead)
    # --------------------------------------------------------

    Xf[iexp, :] = Xprd.isel( time=slice(0, NLEAD)).values

    Yf[iexp, :] = Yprd.isel( time=slice(0, NLEAD)).values

    Zf[iexp, :] = Zprd.isel( time=slice(0, NLEAD)).values

    # --------------------------------------------------------
    # Define verification times
    # --------------------------------------------------------

    verification_times = initial_time + LEAD_TIMES

    # --------------------------------------------------------
    # Store corresponding true states
    #
    # Xtrue(t_initial, t_lead)
    # --------------------------------------------------------

    Xtrue[iexp, :] = Xctl.sel(
        time=verification_times,
        method="nearest",
        tolerance=DT_CTL / 2.0
    ).values

    Ytrue[iexp, :] = Yctl.sel(
        time=verification_times,
        method="nearest",
        tolerance=DT_CTL / 2.0
    ).values

    Ztrue[iexp, :] = Zctl.sel(
        time=verification_times,
        method="nearest",
        tolerance=DT_CTL / 2.0
    ).values

    fin_prd.close()


# ============================================================
# Check the dimensions
# ============================================================

print("========================================")
print("Forecast array dimensions")
print("========================================")

print("Xf shape    :", Xf.shape)
print("Xtrue shape :", Xtrue.shape)

print()
print("Array dimensions are:")
print("(number of experiments, number of lead times)")


# ============================================================
# Create arrays for skill metrics
# ============================================================

ACC_X = np.full(NLEAD, np.nan)
ACC_Y = np.full(NLEAD, np.nan)
ACC_Z = np.full(NLEAD, np.nan)

RMSE_X = np.full(NLEAD, np.nan)
RMSE_Y = np.full(NLEAD, np.nan)
RMSE_Z = np.full(NLEAD, np.nan)

RMSE_CLIM_X = np.full(NLEAD, np.nan)
RMSE_CLIM_Y = np.full(NLEAD, np.nan)
RMSE_CLIM_Z = np.full(NLEAD, np.nan)

RMSESS_X = np.full(NLEAD, np.nan)
RMSESS_Y = np.full(NLEAD, np.nan)
RMSESS_Z = np.full(NLEAD, np.nan)


# ============================================================
# Calculate ACC, RMSE, and RMSE Skill Score
# ============================================================

for ilead in range(NLEAD):

    # --------------------------------------------------------
    # ACC
    # --------------------------------------------------------

    ACC_X[ilead] = np.corrcoef(
        Xtrue[:, ilead],
        Xf[:, ilead]
    )[0, 1]

    ACC_Y[ilead] = np.corrcoef(
        Ytrue[:, ilead],
        Yf[:, ilead]
    )[0, 1]

    ACC_Z[ilead] = np.corrcoef(
        Ztrue[:, ilead],
        Zf[:, ilead]
    )[0, 1]

    # --------------------------------------------------------
    # Forecast RMSE
    # --------------------------------------------------------

    RMSE_X[ilead] = np.sqrt(
        np.mean(
            (Xf[:, ilead] - Xtrue[:, ilead])**2
        )
    )

    RMSE_Y[ilead] = np.sqrt(
        np.mean(
            (Yf[:, ilead] - Ytrue[:, ilead])**2
        )
    )

    RMSE_Z[ilead] = np.sqrt(
        np.mean(
            (Zf[:, ilead] - Ztrue[:, ilead])**2
        )
    )

    # --------------------------------------------------------
    # Climatological-reference RMSE
    # --------------------------------------------------------

    RMSE_CLIM_X[ilead] = np.sqrt(
        np.mean(
            (Xclim - Xtrue[:, ilead])**2
        )
    )

    RMSE_CLIM_Y[ilead] = np.sqrt(
        np.mean(
            (Yclim - Ytrue[:, ilead])**2
        )
    )

    RMSE_CLIM_Z[ilead] = np.sqrt(
        np.mean(
            (Zclim - Ztrue[:, ilead])**2
        )
    )

    # --------------------------------------------------------
    # RMSE Skill Score
    #
    # RMSESS > 0 : forecast is better than climatology
    # RMSESS = 0 : forecast is equal to climatology
    # RMSESS < 0 : forecast is worse than climatology
    # --------------------------------------------------------

    RMSESS_X[ilead] = (
        1.0 - RMSE_X[ilead] / RMSE_CLIM_X[ilead]
    )

    RMSESS_Y[ilead] = (
        1.0 - RMSE_Y[ilead] / RMSE_CLIM_Y[ilead]
    )

    RMSESS_Z[ilead] = (
        1.0 - RMSE_Z[ilead] / RMSE_CLIM_Z[ilead]
    )


# ============================================================
# Print the skill metrics
# ============================================================

print("========================================")
print("Hindcast skill")
print("========================================")

for ilead in range(NLEAD):

    print(f"Lead time = {LEAD_TIMES[ilead]}")

    print(
        f"  X: ACC={ACC_X[ilead]:6.3f}, "
        f"RMSE={RMSE_X[ilead]:6.3f}, "
        f"RMSESS={RMSESS_X[ilead]:6.3f}"
    )

    print(
        f"  Y: ACC={ACC_Y[ilead]:6.3f}, "
        f"RMSE={RMSE_Y[ilead]:6.3f}, "
        f"RMSESS={RMSESS_Y[ilead]:6.3f}"
    )

    print(
        f"  Z: ACC={ACC_Z[ilead]:6.3f}, "
        f"RMSE={RMSE_Z[ilead]:6.3f}, "
        f"RMSESS={RMSESS_Z[ilead]:6.3f}"
    )


# ============================================================
# Create the 3 x 3 figure
# ============================================================

fig, ax = plt.subplots(
    nrows=3,
    ncols=3,
    figsize=(12, 10),
    sharex=True
)

fig.suptitle(
    figure_title,
    fontsize=14
)


# ============================================================
# Row 1: X
# ============================================================

ax[0, 0].plot( LEAD_TIMES, ACC_X, color="C0", linewidth = 2.0)
ax[0, 1].plot( LEAD_TIMES, RMSE_X, color="C0", linewidth = 2.0)
ax[0, 2].plot( LEAD_TIMES, RMSESS_X, color="C0", linewidth = 2.0)

# ============================================================
# Row 2: Y
# ============================================================

ax[1, 0].plot( LEAD_TIMES, ACC_Y, color="C0", linewidth = 2.0)
ax[1, 1].plot( LEAD_TIMES, RMSE_Y, color="C0", linewidth = 2.0)
ax[1, 2].plot( LEAD_TIMES, RMSESS_Y, color="C0", linewidth = 2.0)

# ============================================================
# Row 3: Z
# ============================================================

ax[2, 0].plot( LEAD_TIMES, ACC_Z, color="C0", linewidth = 2.0)
ax[2, 1].plot( LEAD_TIMES, RMSE_Z, color="C0", linewidth = 2.0)
ax[2, 2].plot( LEAD_TIMES, RMSESS_Z, color="C0", linewidth = 2.0)

# ============================================================
# Column titles
# ============================================================

ax[0, 0].set_title("ACC")
ax[0, 1].set_title("RMSE")
ax[0, 2].set_title("RMSE Skill Score")


# ============================================================
# Row labels
# ============================================================

ax[0, 0].set_ylabel("X")
ax[1, 0].set_ylabel("Y")
ax[2, 0].set_ylabel("Z")


# ============================================================
# Reference lines
# ============================================================

for irow in range(3):

    # ACC = 0 indicates no linear correlation
    ax[irow, 0].axhline(
        0.0,
        color="black",
        linestyle="--",
        linewidth=0.8
    )

    # RMSE Skill Score = 0 indicates no improvement
    # relative to climatology
    ax[irow, 2].axhline(
        0.0,
        color="black",
        linestyle="--",
        linewidth=0.8
    )


# ============================================================
# Axis limits
# ============================================================

for irow in range(3):

    ax[irow, 0].set_ylim(-1.0, 1.0)
    ax[irow, 2].set_ylim(-1.0, 1.0)


# ============================================================
# Common plot settings
# ============================================================

XTICKS = np.arange(0, FORECAST_LENGTH + 1, 1)

for irow in range(3):
    for icol in range(3):
        ax[irow, icol].set_xlim(0.0, FORECAST_LENGTH)
        ax[irow, icol].set_xticks(XTICKS)
        ax[irow, icol].grid(True)


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


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

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

plt.show()
plt.close()

print("Figure saved as:", outfile_png)


# ============================================================
# Close the control dataset
# ============================================================

fin_ctl.close()