difficulty_calibration.py

"""
Benchmark Difficulty Calibration Framework.

Constructs a 4-dimension composite difficulty score for multivariate
time series forecasting datasets (D_ns, D_horizon, D_noise, D_dim).
Validates calibration via Spearman correlation against actual MSE.

Reference: Section 3.7 of the accompanying thesis.
"""
import numpy as np
import pandas as pd
from pathlib import Path
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from scipy import stats

OUTPUT_DIR = Path("figures")
OUTPUT_DIR.mkdir(exist_ok=True)

DATASET_NAMES = ["Electricity", "ETTh1", "Exchange", "Traffic", "Weather"]
PRED_HORIZONS = [96, 192, 336, 720]

# Statistical features calibrated in the thesis
STATS = {
    "Electricity": {"F_t": 0.737, "F_s": 0.628, "ADF_p": 0.0, "dim": 321},
    "ETTh1":       {"F_t": 0.969, "F_s": 0.276, "ADF_p": 0.01, "dim": 7},
    "Exchange":    {"F_t": 0.996, "F_s": 0.025, "ADF_p": 0.88, "dim": 8},
    "Traffic":     {"F_t": 0.108, "F_s": 0.216, "ADF_p": 0.0, "dim": 862},
    "Weather":     {"F_t": 0.002, "F_s": 0.003, "ADF_p": 0.0, "dim": 21},
}

# Average MSE across LSTM/DLinear/FITS from thesis experiments
ACTUAL_MSE = {
    ("Electricity", 96): 0.211, ("Electricity", 192): 0.213, ("Electricity", 336): 0.223, ("Electricity", 720): 0.245,
    ("ETTh1", 96): 0.493, ("ETTh1", 192): 0.539, ("ETTh1", 336): 0.625, ("ETTh1", 720): 0.704,
    ("Exchange", 96): 0.554, ("Exchange", 192): 0.729, ("Exchange", 336): 0.936, ("Exchange", 720): 1.218,
    ("Traffic", 96): 0.515, ("Traffic", 192): 0.478, ("Traffic", 336): 0.505, ("Traffic", 720): 0.582,
    ("Weather", 96): 0.336, ("Weather", 192): 0.379, ("Weather", 336): 0.438, ("Weather", 720): 0.560,
}


def compute_difficulty_score(stat_features: dict, horizon: int,
                              w1: float = 0.30, w2: float = 0.25,
                              w3: float = 0.25, w4: float = 0.20) -> dict:
    """
    Compute composite difficulty score from 4 sub-components.

    D_ns  — Non-stationarity penalty (ADF p >= 0.05 triggers)
    D_H   — Horizon scaling (log-normalized against max horizon)
    D_n   — Noise penalty (1 - max(F_t, F_s))
    D_dim — Dimensionality complexity (log-scaled)
    """
    is_ns = 1.0 if stat_features["ADF_p"] >= 0.05 else 0.0
    D_ns = is_ns * (1.0 + np.log(1 + stat_features["ADF_p"]) * 0.5)

    D_horizon = np.log(horizon / 96.0 + 1) / np.log(720 / 96.0 + 1)

    structure = max(stat_features["F_t"], stat_features["F_s"])
    D_noise = max(0, 1.0 - structure)

    D_dim = np.log(1 + stat_features["dim"]) / np.log(1 + 862)

    composite = w1 * D_ns + w2 * D_horizon + w3 * D_noise + w4 * D_dim

    return {
        "composite": round(composite, 4),
        "D_ns": round(D_ns, 4),
        "D_horizon": round(D_horizon, 4),
        "D_noise": round(D_noise, 4),
        "D_dim": round(D_dim, 4),
        "level": "High" if composite > 0.45 else ("Medium" if composite > 0.25 else "Low"),
    }


def calibrate_all():
    """Calibrate difficulty for all 20 (dataset x horizon) combinations."""
    rows = []
    for name in DATASET_NAMES:
        s = STATS[name]
        for h in PRED_HORIZONS:
            score = compute_difficulty_score(s, h)
            actual = ACTUAL_MSE.get((name, h), np.nan)
            rows.append({
                "Dataset": name,
                "Horizon": h,
                **score,
                "Actual_Avg_MSE": round(actual, 4),
            })
    return pd.DataFrame(rows)


def validate_calibration(df: pd.DataFrame):
    """Spearman & Pearson correlation between difficulty score and actual MSE."""
    valid = df[df["Actual_Avg_MSE"].notna()].copy()
    if len(valid) < 5:
        return None, None
    r, p = stats.spearmanr(valid["composite"], valid["Actual_Avg_MSE"])
    pr, pp = stats.pearsonr(valid["composite"], valid["Actual_Avg_MSE"])
    return {"spearman_r": r, "spearman_p": p, "pearson_r": pr, "pearson_p": pp}, valid


def plot_difficulty_vs_mse(df_valid: pd.DataFrame):
    """Scatter plot: difficulty score vs actual MSE with linear fit."""
    fig, ax = plt.subplots(figsize=(8, 6))
    colors = {
        "Electricity": "tab:blue", "ETTh1": "tab:orange",
        "Exchange": "tab:red", "Traffic": "tab:green", "Weather": "tab:purple",
    }
    for name in DATASET_NAMES:
        sub = df_valid[df_valid["Dataset"] == name]
        ax.scatter(sub["composite"], sub["Actual_Avg_MSE"], c=colors[name],
                   label=name, s=80, edgecolors="black", linewidth=0.5)
        for _, row in sub.iterrows():
            ax.annotate(f"H={int(row['Horizon'])}",
                        (row["composite"], row["Actual_Avg_MSE"]),
                        textcoords="offset points", xytext=(5, 5), fontsize=7)

    if len(df_valid) >= 5:
        m, b = np.polyfit(df_valid["composite"], df_valid["Actual_Avg_MSE"], 1)
        x_line = np.linspace(df_valid["composite"].min(), df_valid["composite"].max(), 50)
        ax.plot(x_line, m * x_line + b, "k--", alpha=0.5, label="Linear fit")

    ax.set_xlabel("Difficulty Score (Composite)")
    ax.set_ylabel("Average MSE (3 models)")
    ax.set_title("Benchmark Difficulty Calibration: Score vs Actual MSE")
    ax.legend(fontsize=9)
    ax.grid(True, alpha=0.3)
    fig.tight_layout()
    fig.savefig(OUTPUT_DIR / "difficulty_vs_mse.pdf", dpi=150, bbox_inches="tight")
    plt.close(fig)
    print(f"Figure saved to {OUTPUT_DIR / 'difficulty_vs_mse.pdf'}")


def main():
    df = calibrate_all()
    print("\n=== Difficulty Calibration Results ===")
    print(df[["Dataset", "Horizon", "composite", "level", "Actual_Avg_MSE"]].to_string(index=False))

    corr, df_valid = validate_calibration(df)
    if corr:
        print(f"\n=== Calibration Validation ===")
        print(f"Spearman r = {corr['spearman_r']:.3f} (p = {corr['spearman_p']:.3f})")
        print(f"Pearson r  = {corr['pearson_r']:.3f} (p = {corr['pearson_p']:.3f})")
        plot_difficulty_vs_mse(df_valid)

    # Also write CSV
    df.to_csv("difficulty_scores.csv", index=False)
    print("CSV saved to difficulty_scores.csv")


if __name__ == "__main__":
    main()