#!/usr/bin/env python3

# sera_sert_analysis.py
# Compute SERT (Squared Error–Relevance at Threshold) and SERA (Squared Error–Relevance Area)
# and their normalized SERT-based variants for multiple phi-methods, models, and datasets.

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

# Usage:
#     python sera_sert_analysis.py --dir predictions/ --pattern "*_preds_default_*.csv"

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

def compute_sert_normalized(dataset_name, dt, t0):

    # Style settings
    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
    })

    # --- Argument parsing ---
    parser = argparse.ArgumentParser(
        description="Compute and plot raw vs adjusted vs SERT-normalized SER/SERA metrics"
    )
    parser.add_argument('--dir', dest='data_dir', default='predictions/', help='Directory with prediction CSVs')
    parser.add_argument('--pattern', dest='pattern', default=f'{dataset_name}_preds_default_*.csv', help='Glob pattern for CSV files')
    args = parser.parse_args()

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

    df_list = []
    for path in paths:
        fname = os.path.basename(path)
        parts = fname.split('_')
        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(path)
        # keep raw y_true, y_pred unchanged
        df['dataset'], df['phi_method'], df['model'] = dataset, phi_method, 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 not 0 < 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())
    idx0 = int(np.argmin(np.abs(thresholds - t0)))
    t0 = thresholds[idx0]
    print(f"Using t0 = {t0}")
    # Gaussian weights for adjustment
    sigma = t0/2
    weight = np.exp(-0.5*((thresholds - t0)/sigma)**2)

    # --- SER curve function ---
    def compute_ser(df, true_col, pred_col):
        return np.array([((df[pred_col] - df[true_col])**2)[df['phi'] >= t].sum() for t in thresholds])

    # --- Collect metrics ---
    results = []
    for (dataset, phi_method, model), grp in all_df.groupby(['dataset','phi_method','model']):
        # raw and adjusted SER curves on raw targets
        ser = compute_ser(grp, 'y_true', 'y_pred')
        adj = np.where(thresholds <= t0, ser*weight, ser)

        # SERT-normalized curves (divide by SERT at t=0)
        base = ser[0] if ser[0] > 0 else 1.0
        ser_norm = ser / base
        adj_norm = np.where(thresholds <= t0, ser_norm*weight, ser_norm)

        # areas (SERA)
        raw_area      = np.trapz(ser,       thresholds)
        adj_area      = np.trapz(adj,       thresholds)
        norm_raw_area = np.trapz(ser_norm,  thresholds)
        norm_adj_area = np.trapz(adj_norm,  thresholds)

        results.append({
            'Dataset':         dataset,
            'Model':           model,
            'Method':          phi_method,
            # 'Raw SERA':        raw_area,
            # 'Adj SERA':        adj_area,
            'SERT Normalized Raw SERA':   norm_raw_area,
            'SERT Normalized Adj SERA':   norm_adj_area
        })

    summary = pd.DataFrame(results).sort_values(['Dataset','SERT Normalized Adj SERA'])
    summary.to_csv(f"/Users/bhavneetsingh/Desktop/RiskReg-main/sert&sera_sert_normalized/{dataset_name}_SERT_SERA.csv", index=True)

    # --- Plotting SERT-normalized curves ---
    combos = summary[['Dataset','Model']].drop_duplicates().values.tolist()
    n, cols = len(combos), 2
    rows = int(np.ceil(n/cols))
    fig, axarr = plt.subplots(rows, cols, figsize=(8,3*rows), constrained_layout=True)
    axes = axarr.flatten()
    for i,(dataset,model) in enumerate(combos):
        grp = all_df[(all_df['dataset']==dataset)&(all_df['model']==model)]
        ser = compute_ser(grp, 'y_true', 'y_pred')
        ser_norm = ser / (ser[0] if ser[0]>0 else 1.0)
        adj_norm = np.where(thresholds<=t0, ser_norm*weight, ser_norm)
        ax = axes[i]
        ax.plot(thresholds, ser_norm, '-', marker='o', label='ser_norm')
        ax.plot(thresholds, adj_norm, '--', marker='x', label='adj_norm')
        ax.axvline(t0, linestyle=':', label=f't0={t0}')
        ax.set_title(f"{dataset} - {model} (SERT-normalized)")
        ax.set_xlabel('t'); ax.set_ylabel('Relative SERT')
        ax.set_xticks(thresholds)
        ax.legend(loc='upper right', ncol=2)
    for j in range(n, len(axes)):
        fig.delaxes(axes[j])
    plt.show()

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