
# sera_sert_analysis.py

# Compute SERT (Squared Error–Relevance at Threshold) and SERA (Squared Error–Relevance Area)
# for multiple phi-methods and models from CSV prediction files.

# Formulas:
#   SERTₜ = ∑_{i: φᵢ ≥ t} (y_pred,i - y_true,i)²
#   SERA   = ∑_{i} (y_pred,i - y_true,i)² · φᵢ

# Usage:
#     python sera_sert_analysis.py

# Requirements:
#     pandas, numpy, matplotlib (optional for plotting)

import glob
import os
import argparse
import pandas as pd
import numpy as np
import matplotlib.ticker as mticker
import matplotlib.pyplot as plt

def compute_sert_sera(dataset_name, dt, t0, predictions_dir=None, out_dir=None, show_plot=False):

    # Adjust global font sizes for clarity
    plt.rcParams.update({
        'font.size': 10,
        'axes.titlesize': 12,
        'axes.labelsize': 10,
        'legend.fontsize': 8,
        'xtick.labelsize': 8,
        'ytick.labelsize': 8,
        'lines.linewidth': 1.5,
        'lines.markersize': 5,
        'figure.dpi': 100
    })

    if predictions_dir is None:
        predictions_dir = 'predictions/'
    if out_dir is None:
        out_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', 'sert&sera')

    pattern = f'{dataset_name}_preds_default_*.csv'

    # --- Load data ---
    search_path = os.path.join(predictions_dir, pattern)
    files = glob.glob(search_path)
    if not files:
        raise FileNotFoundError(f"No files matching {search_path}")

    df_list = []
    for fn in files:
        parts = os.path.basename(fn).split('_')
        if len(parts) < 3: continue
        if len(parts) == 5:
            dataset, phi_method, model = parts[0] + parts[1], parts[3], parts[4].replace('.csv','')
        else:
            dataset, phi_method, model = parts[0], parts[2], parts[3].replace('.csv','')
        df = pd.read_csv(fn)
        df['phi_method'] = phi_method
        df['model'] = model
        df_list.append(df)
    all_df = pd.concat(df_list, ignore_index=True)

    # --- User inputs ---
    # dt = float(input("Enter Δt (0<Δt≤1, e.g. 0.05): ").strip())
    if dt <= 0 or dt > 1:
        raise ValueError("Δt must be between 0 and 1.")
    thresholds = np.round(np.arange(0, 1+dt/2, dt), 3)
    print("Thresholds:", thresholds)
    # t0 = float(input("Enter cutoff t0 (0≤t0≤1): ").strip())
    t0_idx = int(np.argmin(np.abs(thresholds - t0)))
    t0 = thresholds[t0_idx]
    print(f"Using t0 = {t0}")

    # --- Precompute Gaussian PDF centered at t0 ---
    sigma = t0 / 2
    weight = np.exp(-0.5 * ((thresholds - t0)/sigma)**2)

    # --- Function for SER curve ---
    def compute_ser_curve(df):
        return np.array([
            ((df['y_pred'] - df['y_true'])**2)[df['phi'] >= t].sum()
            for t in thresholds
        ])

    # --- Collect and summarize results ---
    results = []
    for (phi_method, model), grp in all_df.groupby(['phi_method','model']):
        ser = compute_ser_curve(grp)
        adj = np.where(thresholds <= t0, ser * weight, ser)
        raw_area = np.trapz(ser, thresholds)
        adj_area = np.trapz(adj, thresholds)
        results.append({'Dataset': dataset, 'Model': model, 'Method': phi_method,
                        'Raw SERA': raw_area, 'Adj SERA': adj_area})
    summary_df = pd.DataFrame(results).sort_values('Adj SERA')
    os.makedirs(out_dir, exist_ok=True)
    summary_df.to_csv(os.path.join(out_dir, f"{dataset_name}_SERT_SERA.csv"), index=True)

    # --- Plotting with tight spacing ---
    models = summary_df['Model'].unique()
    n = len(models)
    cols = 2
    rows = int(np.ceil(n/cols))
    fig, axes = plt.subplots(rows, cols, figsize=(8, 3*rows), constrained_layout=True)
    axes = axes.flatten()

    for i, model in enumerate(models):
        ax = axes[i]
        for phi_method, grp in all_df[all_df['model']==model].groupby('phi_method'):
            ser = compute_ser_curve(grp)
            adj = np.where(thresholds <= t0, ser * weight, ser)
            ax.plot(thresholds, ser, '-', marker='o', label='raw')
            ax.plot(thresholds, adj, '--', marker='x', label='adj')
        ax.axvline(t0, color='red', linestyle=':', label=f't0={t0}')
        ax.set_title(model)
        ax.set_xlabel('t')
        ax.set_ylabel('SER')
        ax.set_xlim(0, 1)  # keep range [0,1]
        ax.xaxis.set_major_locator(mticker.MultipleLocator(0.1))   # ticks every 0.10
        ax.xaxis.set_minor_locator(mticker.MultipleLocator(0.05))  # (optional) minor ticks
        ax.xaxis.set_major_formatter(mticker.FormatStrFormatter('%.2f'))  # show 0.00, 0.10, ...
        ax.legend(loc='upper right', ncol=2)

    # Remove unused axes if any
    total_plots = rows*cols
    for j in range(n, total_plots):
        fig.delaxes(axes[j])

    plot_path = os.path.join(out_dir, f"{dataset_name}_SER_curves.png")
    fig.savefig(plot_path, bbox_inches="tight")
    if show_plot:
        plt.show()
    plt.close(fig)

    return {
        "csv_path": os.path.join(out_dir, f"{dataset_name}_SERT_SERA.csv"),
        "plot_path": plot_path,
        "summary": summary_df.to_dict(orient="records"),
    }


if __name__ == "__main__":
    for datasets in os.listdir("data"):
        dataset_name = datasets.split(".")[0]
        if dataset_name == "":
            continue
        compute_sert_sera(dataset_name, 0.05, 0.5)