# sera_sert_analysis.py
# Raw SERT/SERA + Normalized SERT/SERA (unitless, cross-dataset)
# Compute SERT (Squared Error–Relevance at Threshold) and SERA (Squared Error–Relevance Area)
# for multiple phi-methods and models from CSV prediction files.
#
# Raw (existing):
#   SERT_t = ∑_{i: φ_i ≥ t} (y_pred,i − y_true,i)^2
#   SERA   = ∫_0^1 SERT_t dt           (numerical: trapz over thresholds)
#   Adjusted SERA uses a Gaussian weight w(t) applied to SERT_t for t ≤ t0.
#
# Normalized (unitless & cross-dataset comparable):
#   Let N be total rows and s_y a dataset scale of y_true (std/iqr/max/mean via --scale_mode).
#   Define cumulative normalized SERT two ways (choose via --norm_style):
#     NSSE_t  = SER_t / (N · s_y^2)                 (--norm_style sse)
#     NRMSE_t = sqrt( SER_t / (N · s_y^2) )         (--norm_style rmse, default)
#   Adjusted normalized uses the same Gaussian on SER_t first, then normalizes:
#     NSSE_adj_t  = SER_adj_t / (N · s_y^2)
#     NRMSE_adj_t = sqrt( SER_adj_t / (N · s_y^2) )
#   Areas (reported in the table; all unitless):
#     Norm SERA      = ∫_0^1 NSSE_t  dt  or ∫_0^1 NRMSE_t  dt
#     Norm Adj SERA  = ∫_0^1 NSSE_adj_t dt or ∫_0^1 NRMSE_adj_t dt
#
# Threshold grid:
#   You enter Δt and t0; the code builds thresholds on [0,1] with step Δt and
#   force-inserts t0 exactly (no “0.55 when you choose 0.50” issue).
#
# Output table (per dataset × phi_method × model):
#   Dataset, Method, Model,
#   Raw SERA, Adj SERA,
#   Norm SERA, Norm Adj SERA,
#   ScaleMode (std/iqr/max/mean), NormStyle (rmse/sse).
#
# Plot:
#   ONE figure with separate subplots per model.
#   Each subplot shows ONLY the normalized curves (unitless):
#     {NRMSE or NSSE}_t and its adjusted version {…}_adj_t for each phi_method.
#   X-axis ticks are thinned via --xtick_step for readability; a red line marks t0.
#
# Usage (examples):
#   python sera_sert_analysis.py --dir predictions --pattern "College_preds_default_*.csv" \
#       --scale_mode std --norm_style rmse
#   python sera_sert_analysis.py --dir predictions --pattern "housing_preds_default_*.csv" \
#       --scale_mode iqr --norm_style sse
#
# Requirements:
#   pandas, numpy, matplotlib

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

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

    # --------- Args ----------
    parser = argparse.ArgumentParser(
        description="Compute raw/adjusted SERA and normalized SERA; plot normalized curves per model."
    )
    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')
    parser.add_argument('--scale_mode', choices=['std','iqr','max','mean'], default='std',
                        help='Dataset scale s_y used in normalization (default: std)')
    parser.add_argument('--norm_style', choices=['rmse','sse'], default='rmse',
                        help='rmse: sqrt(SER/(N*s_y^2)); sse: SER/(N*s_y^2) (no sqrt).')
    parser.add_argument("--xtick_step", type=int, default=2,
                        help='Plot every k-th tick on x-axis for readability.')
    args = parser.parse_args()

    # --------- Load 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:
        base = os.path.basename(fn)
        parts = base.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(fn)
        # Expect columns: y_true, y_pred, phi
        if not {'y_true','y_pred','phi'}.issubset(df.columns):
            raise ValueError(f"Missing one of required columns (y_true,y_pred,phi) in {fn}")
        df['dataset'] = dataset
        df['phi_method'] = phi_method
        df['model'] = model
        df_list.append(df)

    all_df = pd.concat(df_list, ignore_index=True)

    # --------- Interactive 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.")

    # base grid then force-include t0 exactly
    base_thresholds = np.round(np.arange(0, 1 + dt/2, dt), 6)
    print("Base thresholds:", base_thresholds)

    # t0_user = float(input("Enter cutoff t0 (0≤t0≤1): ").strip())
    t0_user = np.clip(t0, 0.0, 1.0)

    thresholds = np.unique(np.sort(np.append(base_thresholds, t0_user)))
    t0_idx = int(np.where(np.isclose(thresholds, t0_user))[0][0])
    t0 = float(thresholds[t0_idx])
    print(f"Using t0 = {t0:.2f}")

    # --------- Gaussian weights (for raw adjusted SERA) ----------
    sigma = t0 / 2 if t0 > 0 else max(dt, 1e-6)  # avoid sigma=0
    weight = np.exp(-0.5 * ((thresholds - t0) / max(sigma, 1e-12))**2)

    # --------- Raw SER curve (for raw/adjusted SERA) ----------
    def compute_ser_curve(df, thresholds):
        se = (df['y_pred'] - df['y_true'])**2
        phi = df['phi']
        return np.array([se[phi >= t].sum() for t in thresholds], dtype=float)

    # --------- Helpers for normalized SERT ----------
    def _dataset_scale(y, mode='std'):
        y = np.asarray(y, float)
        if mode == 'std':  return max(np.std(y, ddof=1), 1e-12)
        if mode == 'iqr':
            q3, q1 = np.percentile(y, [75, 25]); return max(q3 - q1, 1e-12)
        if mode == 'max':  return max(np.max(y), 1e-12)
        if mode == 'mean': return max(np.mean(y), 1e-12)
        return 1.0

    def compute_norm_sert_curve(df, thresholds, scale_mode='std', norm_style='rmse'):
        """
        Normalized cumulative SERT:
        NSSE_t = SER_t / (N * s_y^2)         # norm_style='sse'
        NRMSE_t = sqrt(NSSE_t)               # norm_style='rmse'
        Adjusted version: apply Gaussian to SER_t first, then normalize.
        """
        y   = df['y_true'].to_numpy(float)
        yp  = df['y_pred'].to_numpy(float)
        phi = df['phi'].to_numpy(float)
        N   = len(y)

        s_y = _dataset_scale(y, scale_mode)
        se  = (yp - y)**2

        # cumulative SSE → monotone
        ser = np.array([se[phi >= t].sum() for t in thresholds], dtype=float)
        # adjusted cumulative SSE (same logic as raw)
        adj_ser = np.where(thresholds <= t0, ser * weight, ser)

        denom = max(N * (s_y**2), 1e-12)
        nsse_curve = ser / denom
        nsse_adj   = adj_ser / denom

        if norm_style == 'rmse':
            curve = np.sqrt(nsse_curve)
            adj   = np.sqrt(nsse_adj)
        else:  # 'sse'
            curve = nsse_curve
            adj   = nsse_adj

        area_raw = float(np.trapz(curve, thresholds))
        area_adj = float(np.trapz(adj,   thresholds))
        return curve, adj, area_raw, area_adj

    # --------- Compute table rows ----------
    results = []
    group_keys = ['dataset', 'phi_method', 'model']
    for (dataset, phi_method, model), grp in all_df.groupby(group_keys):
        # Raw/Adjusted SERA (originals)
        ser = compute_ser_curve(grp, thresholds)
        ser_adj = np.where(thresholds <= t0, ser * weight, ser)
        raw_sera = float(np.trapz(ser, thresholds))
        adj_sera = float(np.trapz(ser_adj, thresholds))

        # Normalized SERA (RMSE-style or SSE-style)
        curve, adj_curve, norm_sera, norm_adj_sera = compute_norm_sert_curve(
            grp, thresholds, args.scale_mode, args.norm_style
        )

        results.append({
            'Dataset': dataset, 'Method': phi_method, 'Model': model, 
            # 'NormStyle': args.norm_style, 'ScaleMode': args.scale_mode, 
            # 'Raw SERA': raw_sera, 'Adj SERA': adj_sera,
            'NRMSE Normalized Raw SERA': norm_sera, 'NRMSE Normalized Adj SERA': norm_adj_sera
        })

    summary_df = pd.DataFrame(results).sort_values(['Dataset', 'Method', 'Model'])
    summary_df.to_csv(f"/Users/bhavneetsingh/Desktop/RiskReg-main/sert&sera_nrmse_normalized/{dataset_name}_SERT_SERA.csv", index=True)

    # --- Plot: ONE figure, separate subplot per MODEL (normalized only) ---
    models = sorted(all_df['model'].unique())
    n = len(models)
    cols = min(3, n)
    rows = int(np.ceil(n / cols))

    fig, axes = plt.subplots(rows, cols, figsize=(5.8*cols, 4.0*rows), constrained_layout=True)
    axes = np.atleast_1d(axes).ravel()

    linestyles = ['-', '--', ':', '-.']  # one linestyle per phi_method
    style_map = {}
    y_label = ('NRMSE-style' if args.norm_style == 'rmse' else 'NSSE-style') + ' SERT (unitless)'

    for i, model in enumerate(models):
        ax = axes[i]
        for phi_method, grp in all_df[all_df['model'] == model].groupby('phi_method'):
            curve, adj_curve, _, _ = compute_norm_sert_curve(
                grp, thresholds, args.scale_mode, args.norm_style
            )
            if phi_method not in style_map:
                style_map[phi_method] = linestyles[len(style_map) % len(linestyles)]
            ls = style_map[phi_method]
            ax.plot(thresholds, curve, linestyle=ls, marker='o', alpha=0.95,
                    label=f'{phi_method} {args.norm_style.upper()}_t')
            ax.plot(thresholds, adj_curve, linestyle=ls, marker='x', alpha=0.95,
                    label=f'{phi_method} {args.norm_style.upper()}_adj_t')

        ax.axvline(t0, color='red', linestyle=':', linewidth=1.2, label='t0')
        ax.set_title(model)
        ax.set_xlabel('t')
        ax.set_ylabel(y_label)

        # cleaner x ticks
        step = max(1, args.xtick_step)
        tick_pos = thresholds[::step]
        ax.set_xticks(tick_pos)
        ax.set_xticklabels([f"{x:.2f}" for x in tick_pos])
        ax.grid(alpha=0.25, linestyle=":", linewidth=0.7)
        ax.legend(loc='upper right', fontsize=8, ncol=1)

    # remove unused axes
    for j in range(i+1, 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_nrmse_nsse_normalized(dataset_name, 0.05, 0.5)