"""
09 — Dataset size sensitivity study + error decomposition.

Trains a fresh POD-NN at each of several training sizes against a
fixed 400-sample held-out test set. For each model, total Test MSE
is decomposed into MSE_POD (truncation error: Y_test vs POD-projected
Y_test) and MSE_NN (regression error: NN-predicted vs POD-projected).

Outputs:
  - dataset_size_sensitivity.png — Test MSE & Max AE vs N (paper Fig 3a)
  - error_decomposition.png      — MSE_POD vs MSE_NN per N (paper Fig 5)
  - dataset_size_sensitivity.csv — raw numbers
"""

import csv
import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import mean_squared_error

from _utils import (
    load_data, train_pod_nn, predict_field, seed_all,
    TEST_SIZE, N_MODES, EPOCHS, ANALYZED_DIR, PLOTS_DIR,
)

SIZES = [10, 25, 50, 75, 100, 150, 200, 500, 1000, 1600]


def main():
    parameters, fields, _, _ = load_data()
    X_test = parameters[-TEST_SIZE:]
    Y_test = fields[-TEST_SIZE:]
    n_pool = len(parameters) - TEST_SIZE

    results = []
    for size in SIZES:
        if size > n_pool:
            print(f"Skipping size={size} (exceeds train pool {n_pool})")
            continue
        seed_all()
        print(f"\n=== Training N={size} ===")
        trained = train_pod_nn(parameters[:size], fields[:size],
                               n_modes=N_MODES, epochs=EPOCHS, verbose=0)
        Y_pred = predict_field(trained, X_test)

        mse = mean_squared_error(Y_test.flatten(), Y_pred.flatten())
        max_ae = float(np.max(np.abs(Y_test - Y_pred)))

        # Error decomposition
        pca = trained["pca"]
        Y_pod = pca.inverse_transform(pca.transform(Y_test))
        mse_pod = float(np.mean((Y_pod - Y_test) ** 2))
        mse_nn  = float(np.mean((Y_pred - Y_pod) ** 2))

        results.append({
            "dataset_size": size, "mse": mse, "max_ae": max_ae,
            "mse_pod": mse_pod, "mse_nn": mse_nn,
        })
        print(f"  MSE={mse:.4f}  MaxAE={max_ae:.4f}  "
              f"MSE_POD={mse_pod:.4f}  MSE_NN={mse_nn:.4f}")

    # Save raw numbers
    csv_path = ANALYZED_DIR / "dataset_size_sensitivity.csv"
    csv_path.parent.mkdir(parents=True, exist_ok=True)
    with open(csv_path, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=["dataset_size", "mse", "max_ae",
                                          "mse_pod", "mse_nn"])
        w.writeheader(); w.writerows(results)
    print(f"\nSaved {csv_path}")

    sizes   = [r["dataset_size"] for r in results]
    mses    = [r["mse"]          for r in results]
    max_aes = [r["max_ae"]       for r in results]

    FS = 24
    fig, ax1 = plt.subplots(figsize=(10, 6))
    ax2 = ax1.twinx()
    ax1.plot(sizes, mses,    "o-", color="tab:blue", linewidth=2,
             markersize=8, label=r"Test MSE (K$^2$)")
    ax2.plot(sizes, max_aes, "D-", color="tab:red",  linewidth=2,
             markersize=8, label="Test Max AE (K)")
    ax1.set_xlabel("Training Dataset Size", fontsize=FS)
    ax1.set_ylabel(r"Test MSE (K$^2$)", fontsize=FS, color="tab:blue")
    ax2.set_ylabel("Test Max AE (K)", fontsize=FS, color="tab:red")
    ax1.tick_params(labelsize=FS - 4, axis="both")
    ax2.tick_params(labelsize=FS - 4, axis="y", labelcolor="tab:red")
    ax1.tick_params(axis="y", labelcolor="tab:blue")
    ax1.grid(True, alpha=0.3)

    plt.tight_layout()
    PLOTS_DIR.mkdir(parents=True, exist_ok=True)
    out = PLOTS_DIR / "dataset_size_sensitivity.png"
    plt.savefig(out, dpi=300, bbox_inches="tight")
    print(f"Saved {out}")
    plt.close()

    # --- Error decomposition bar plot (paper Fig 5) ---
    mse_pods = [r["mse_pod"] for r in results]
    mse_nns  = [r["mse_nn"]  for r in results]

    FS2 = 20
    x = np.arange(len(sizes))
    width = 0.35

    fig, ax = plt.subplots(figsize=(14, 6))
    bars_pod = ax.bar(x - width / 2, mse_pods, width,
                      label=r"$\mathrm{MSE_{POD}}$", color="tab:orange",
                      edgecolor="black", linewidth=0.5)
    bars_nn  = ax.bar(x + width / 2, mse_nns, width,
                      label=r"$\mathrm{MSE_{NN}}$", color="tab:purple",
                      edgecolor="black", linewidth=0.5)
    for bar in list(bars_pod) + list(bars_nn):
        ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height(),
                f"{bar.get_height():.1f}", ha="center", va="bottom",
                fontsize=FS2 - 4)
    ax.set_xlabel("Training Dataset Size (N)", fontsize=FS2)
    ax.set_ylabel(r"MSE (K$^2$)", fontsize=FS2)
    ax.set_xticks(x)
    ax.set_xticklabels([str(s) for s in sizes])
    ax.tick_params(labelsize=FS2 - 4)
    ax.legend(fontsize=FS2 - 4)
    ax.grid(True, alpha=0.3, axis="y")

    plt.tight_layout()
    out_decomp = PLOTS_DIR / "error_decomposition.png"
    plt.savefig(out_decomp, dpi=300, bbox_inches="tight")
    print(f"Saved {out_decomp}")
    plt.close()


if __name__ == "__main__":
    main()