#!/usr/bin/env python3

# 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 --dir predictions/ --pattern "preds_default_*.csv"

# Requirements:
#     pandas, numpy, matplotlib

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

def compute_target_normalized(dataset_name, dt, t0):
    # 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
    })

    # --- Argument parsing ---
    parser = argparse.ArgumentParser(
        description="Compute and plot raw vs Gaussian-adjusted SER curves with dataset labels"
    )
    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 and normalize data ---
    search_path = os.path.join(args.data_dir, args.pattern)
    files = glob.glob(search_path)
    if not files:
        raise FileNotFoundError(f"No files matching {search_path}")

    df_list = []
    for fn in files:
        basename = os.path.basename(fn)
        parts = basename.split('_')
        # if len(parts) < 4:
        #     continue
        # dataset = parts[0]
        # phi_method = parts[2]
        # model = parts[3].replace('.csv','')

        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','')

        # Load predictions
        df = pd.read_csv(fn)

        # Min–Max normalize (and clip) y_true and y_pred based on y_true range
        y_min, y_max = df['y_true'].min(), df['y_true'].max()
        if y_max > y_min:
            df['y_true_norm'] = np.clip((df['y_true'] - y_min) / (y_max - y_min), 0.0, 1.0)
            df['y_pred_norm'] = np.clip((df['y_pred'] - y_min) / (y_max - y_min), 0.0, 1.0)
        else:
            # If constant target, set zeros
            df['y_true_norm'] = 0.0
            df['y_pred_norm'] = 0.0

        # Attach metadata
        df['dataset'] = dataset
        df['phi_method'] = phi_method
        df['model'] = model
        df_list.append(df)

    # Combine all data into one DataFrame
    all_df = pd.concat(df_list, ignore_index=True)

    # --- User inputs for thresholds ---
    # 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_norm'] - df['y_true_norm']) ** 2)[df['phi'] >= t].sum()
            for t in thresholds
        ])

    # --- Collect and summarize results ---
    results = []
    for (dataset, phi_method, model), grp in all_df.groupby(['dataset','phi_method','model']):
        ser = compute_ser_curve(grp)
        # Gaussian-adjusted: apply weighting up to t0
        adj = np.where(thresholds <= t0, ser * weight, ser)

        # Compute areas under the curves
        raw_area = np.trapz(ser, thresholds)
        adj_area = np.trapz(adj, thresholds)

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

    # Build summary table
    summary_df = pd.DataFrame(results).sort_values(['Dataset','Target Normalized Adj SERA'])
    print(summary_df.to_string(index=False))
    summary_df.to_csv(f"/Users/bhavneetsingh/Desktop/RiskReg-main/sert&sera_target_normalized/{dataset_name}_SERT_SERA.csv", index=True)

    # --- Plotting with tight spacing ---
    combos = summary_df[['Dataset','Model']].drop_duplicates().values.tolist()
    n = len(combos)
    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, (dataset, model) in enumerate(combos):
        ax = axes[i]
        grp = all_df[(all_df['dataset'] == dataset) & (all_df['model'] == model)]
        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, linestyle=':', label=f't0={t0}')
        ax.set_title(f"{dataset} - {model}")
        ax.set_xlabel('t')
        ax.set_ylabel('SER')
        ax.set_xticks(thresholds)
        ax.legend(loc='upper right', ncol=2)

    # Remove any unused axes
    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_target_normalized(dataset_name, 0.05, 0.5)