"""
BEWARE Results Analysis


Description
-----------
Loads completed nested cross-validation results,
computes performance statistics,
and generates publication-quality figures.

Outputs include

- Cross-validation statistics
- Predicted vs observed scatter plots
- Residual analysis
- Hyperparameter distributions
- Learning curves
"""

# ----------------------------------------


from pathlib import Path
import json

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

plt.rcParams["figure.dpi"] = 300
plt.rcParams["savefig.dpi"] = 300
plt.rcParams["font.size"] = 11

# --------------
# Configuration
# --------------

BASE_DIR = Path(__file__).resolve().parent

CONFIG = {

    "results_folder": BASE_DIR / "Nested_CV_Results",


    "figures_folder": BASE_DIR / "Nested_CV_Results" / "Figures",

}

CONFIG["figures_folder"].mkdir(
    parents=True,
    exist_ok=True
)

print("=" * 50)
print("BEWARE Results Analysis")
print("=" * 50)

print(f"Results folder : {CONFIG['results_folder']}")
print(f"Figures folder : {CONFIG['figures_folder']}")

# ----------------------
# Locate completed folds
# ----------------------

print("\nSearching for completed folds...")

fold_folders = []

for folder in sorted(CONFIG["results_folder"].glob("fold_*")):

    if not folder.is_dir():
        continue

    required_files = [

        folder / "predictions.csv",
        folder / "best_parameters.json",
        folder / "all_trials.csv"

    ]

    if all(f.exists() for f in required_files):

        fold_folders.append(folder)

        print(f"✓ {folder.name}")

    else:

        print(f"✗ {folder.name} (incomplete - skipped)")

print(f"\nCompleted folds found: {len(fold_folders)}")

all_predictions = []
all_trials = []
best_parameters = []

# ---------------------------
# Load results from each fold
# ---------------------------

for fold_folder in fold_folders:

    print(f"\nLoading {fold_folder.name}...")

    #
    # Predictions
    #

    prediction_file = fold_folder / "predictions.csv"

    if prediction_file.exists():

        df = pd.read_csv(prediction_file)
        df.columns = df.columns.str.lower()

        df["fold"] = fold_folder.name

        all_predictions.append(df)

    #
    # Hyperparameter trials
    #

    trial_file = fold_folder / "all_trials.csv"

    if trial_file.exists():

        all_trials.append(
            pd.read_csv(trial_file)
        )

    #
    # Best parameters
    #

    parameter_file = fold_folder / "best_parameters.json"

    if parameter_file.exists():

        with open(parameter_file) as f:

            params = json.load(f)

        params["fold"] = fold_folder.name

        best_parameters.append(params)


# -----------------
# Combine datasets
# -----------------

if all_predictions:

    predictions = pd.concat(
        all_predictions,
        ignore_index=True
    )

else:

    predictions = pd.DataFrame()

if all_trials:

    trials = pd.concat(
        all_trials,
        ignore_index=True
    )

else:

    trials = pd.DataFrame()

if len(best_parameters):

    best_parameters = pd.DataFrame(best_parameters)

else:

    raise RuntimeError("No completed folds found.")


print("\nSummary")

print("-" * 40)

print(f"Predictions : {len(predictions):,}")

print(f"Trials      : {len(trials):,}")

print(f"Best models : {len(best_parameters)}")


from sklearn.metrics import (
    mean_squared_error,
    mean_absolute_error,
    r2_score,
)

# -----------------------------
# Calculate performance metrics
# -----------------------------

from scipy.stats import pearsonr

metrics = []

for fold_name, fold_df in predictions.groupby("fold"):

    y_true = fold_df["observed"].to_numpy()
    y_pred = fold_df["predicted"].to_numpy()

    rmse = np.sqrt(mean_squared_error(y_true, y_pred))
    mae = mean_absolute_error(y_true, y_pred)
    r2 = r2_score(y_true, y_pred)

    r, _ = pearsonr(y_true, y_pred)

    bias = np.mean(y_pred - y_true)

    relative_bias = bias / np.mean(y_true)

    scatter_index = rmse / np.mean(y_true)

    nse = 1 - (
        np.sum((y_true - y_pred) ** 2)
        /
        np.sum((y_true - np.mean(y_true)) ** 2)
    )

    mape = np.mean(
        np.abs((y_true - y_pred) / y_true)
    ) * 100

    willmott = 1 - (
        np.sum((y_pred - y_true) ** 2)
        /
        np.sum(
            (
                np.abs(y_pred - np.mean(y_true))
                +
                np.abs(y_true - np.mean(y_true))
            ) ** 2
        )
    )

    metrics.append({

        "Fold": fold_name,

        "RMSE": rmse,

        "MAE": mae,

        "R2": r2,

        "Pearson_r": r,

        "Bias": bias,

        "Relative_Bias": relative_bias,

        "Scatter_Index": scatter_index,

        "NSE": nse,

        "MAPE": mape,

        "Willmott_d": willmott

    })

metrics_df = pd.DataFrame(metrics)


summary = pd.DataFrame({

    "Metric": [

        "RMSE",
        "MAE",
        "R2",
        "Pearson_r",
        "Bias",
        "Relative_Bias",
        "Scatter_Index",
        "NSE",
        "MAPE",
        "Willmott_d"

    ],

    "Mean": [

        metrics_df["RMSE"].mean(),
        metrics_df["MAE"].mean(),
        metrics_df["R2"].mean(),
        metrics_df["Pearson_r"].mean(),
        metrics_df["Bias"].mean(),
        metrics_df["Relative_Bias"].mean(),
        metrics_df["Scatter_Index"].mean(),
        metrics_df["NSE"].mean(),
        metrics_df["MAPE"].mean(),
        metrics_df["Willmott_d"].mean()

    ],

    "SD": [

        metrics_df["RMSE"].std(),
        metrics_df["MAE"].std(),
        metrics_df["R2"].std(),
        metrics_df["Pearson_r"].std(),
        metrics_df["Bias"].std(),
        metrics_df["Relative_Bias"].std(),
        metrics_df["Scatter_Index"].std(),
        metrics_df["NSE"].std(),
        metrics_df["MAPE"].std(),
        metrics_df["Willmott_d"].std()

    ]

})
print("\n")
print("=" * 50)
print("Overall Cross-Validation Performance")
print("=" * 50)

print(metrics_df.round(5))

print("\nSummary")

print(summary.round(5))

metrics_df.to_csv(
    CONFIG["results_folder"] / "overall_metrics.csv",
    index=False
)

summary.to_csv(
    CONFIG["results_folder"] / "summary_statistics.csv",
    index=False
)

from scipy.stats import pearsonr

r, p = pearsonr(
    predictions["observed"],
    predictions["predicted"]
)

print(f"\nPearson r : {r:.4f}")
print(f"P-value   : {p:.3e}")


# --------------------
# Calculate residuals
# --------------------

predictions["Residual"] = (
    predictions["observed"]
    - predictions["predicted"]
)

# --------------
# Residual Plot
# --------------

print("Generating residual plot...")

plt.figure(figsize=(6,5))

plt.scatter(
    predictions["predicted"],
    predictions["Residual"],
    alpha=0.5,
    s=18
)

plt.axhline(
    0,
    linestyle="--",
    linewidth=2
)

plt.xlabel("Predicted $R^2$")
plt.ylabel("Residual")

plt.title("Residual Plot")

plt.tight_layout()

plt.savefig(
    CONFIG["figures_folder"] / "residual_plot.png"
)

plt.close()


# ----------------------
# Predicted vs Observed
# ----------------------

print("Generating observed vs predicted plot...")

plt.figure(figsize=(6,6))

for fold_name, fold_df in predictions.groupby("fold"):

    plt.scatter(
        fold_df["observed"],
        fold_df["predicted"],
        s=20,
        alpha=0.6,
        label=fold_name
    )

minimum = min(predictions["observed"].min(),
              predictions["predicted"].min())

maximum = max(predictions["observed"].max(),
              predictions["predicted"].max())

plt.plot(
    [minimum, maximum],
    [minimum, maximum],
    "k--",
    linewidth=2,
    label="1:1 line"
)

plt.xlabel("Observed $R^2$")
plt.ylabel("Predicted $R^2$")

plt.title("Observed vs Predicted")

text = (
    f"RMSE = {metrics_df['RMSE'].mean():.3f}\n"
    f"MAE = {metrics_df['MAE'].mean():.3f}\n"
    f"$R^2$ = {metrics_df['R2'].mean():.3f}"
)

plt.text(
    0.05,
    0.95,
    text,
    transform=plt.gca().transAxes,
    verticalalignment="top",
    bbox=dict(facecolor="white", alpha=0.8)
)

plt.legend(title="Fold")

plt.tight_layout()

plt.savefig(
    CONFIG["figures_folder"] / "observed_vs_predicted.png"
)

plt.close()

# -------------------
# Residual Histogram
# -------------------

print("Generating residual histogram...")

plt.figure(figsize=(6,5))

plt.hist(
    predictions["Residual"],
    bins=40
)

plt.xlabel("Residual")

plt.ylabel("Frequency")

plt.title("Residual Distribution")

plt.tight_layout()

plt.savefig(
    CONFIG["figures_folder"] / "residual_histogram.png"
)

plt.close()

# --------------------------
# Hyperparameter distributions
# --------------------------

print("Generating hyperparameter summary...")

numeric = [

    "hidden_layers",
    "neurons",
    "batch_size",
    "dropout",
    "learning_rate",
    "weight_decay",
    "value"

]

available = [c for c in numeric if c in trials.columns]

trials[available].hist(

    figsize=(10,8),

    bins=20

)

plt.tight_layout()

plt.savefig(

    CONFIG["figures_folder"] / "hyperparameter_distributions.png"

)

plt.close()

# -------------------------------------
# Learning Curves (Best Trial Per Fold)
# -------------------------------------

print("Generating learning curves...")

for fold_folder in fold_folders:

    print(f"  {fold_folder.name}")

    # Load all trial results
    trials = pd.read_csv(fold_folder / "all_trials.csv")

    # Best trial = lowest validation loss
    best_trial = trials.loc[trials["value"].idxmin()]

    trial_number = int(best_trial["number"])

    history_file = (
        fold_folder /
        f"trial_{trial_number:03d}_history.csv"
    )

    if not history_file.exists():
        print(f"    Missing history for trial {trial_number}")
        continue

    history = pd.read_csv(history_file)

    plt.figure(figsize=(6,4))

    plt.plot(
        history["epoch"],
        history["train_loss"],
        label="Training",
        linewidth=2
    )

    plt.plot(
        history["epoch"],
        history["validation_loss"],
        label="Validation",
        linewidth=2
    )

    plt.xlabel("Epoch")
    plt.ylabel("Loss")

    plt.title(f"{fold_folder.name} Learning Curve for Best Model")

    plt.legend()

    plt.tight_layout()

    plt.savefig(
        CONFIG["figures_folder"] /
        f"{fold_folder.name}_learning_curve.png"
    )

    plt.close()