src/ablation_study.py

"""
Ablation Studies Module
========================
Evaluates the impact of:
1. Feature engineering strategies (which feature groups matter)
2. Spatial encoding methods (spatial lag, Fourier, graph, clusters)
3. Clustering granularity (K=5,10,20,50 for geographic clusters)
4. Transformer embeddings vs. pure GBDT
5. Calibration methods (Platt vs. Isotonic vs. Temperature vs. Ensemble)
"""

import numpy as np
import pandas as pd
from sklearn.model_selection import StratifiedKFold
from sklearn.metrics import roc_auc_score
import json
from copy import deepcopy

from hybrid_model import (
    HybridModel, evaluate_model, prepare_data, 
    FEATURE_GROUPS, get_feature_columns,
    train_xgboost, train_lightgbm
)
from calibration import (
    PlattScaling, TemperatureScaling, IsotonicCalibration,
    EnsembleCalibration, expected_calibration_error, calibration_report
)
from sklearn.preprocessing import QuantileTransformer


def run_ablation_feature_groups(df, target_col='is_hotspot', n_splits=3):
    """Ablation 1: Impact of each feature group.
    
    Tests: all features, drop-one-group, only-one-group.
    """
    print("\n" + "="*80)
    print("ABLATION 1: Feature Group Impact Analysis")
    print("="*80)
    
    all_feature_cols = get_feature_columns(df, include_spatial=False, include_fourier=False, 
                                           include_graph=False, include_clusters=False)
    
    X_all = df[all_feature_cols].values
    y = df[target_col].values.astype(int)
    
    qt = QuantileTransformer(output_distribution='normal', random_state=42)
    
    results = {}
    
    # Baseline: all feature groups
    print("\n--- All Base Features ---")
    metrics = _cross_validate(X_all, y, qt, n_splits)
    results['all_features'] = metrics
    print(f"  AUC: {metrics['auc_roc']:.4f} | F1: {metrics['f1_macro']:.4f} | Acc: {metrics['accuracy']:.4f}")
    
    # Drop-one-group analysis
    print("\n--- Drop-One-Group Analysis ---")
    for group_name, group_cols in FEATURE_GROUPS.items():
        remaining_cols = [c for c in all_feature_cols if c not in group_cols]
        if len(remaining_cols) == 0:
            continue
        
        X_drop = df[remaining_cols].values
        metrics = _cross_validate(X_drop, y, qt, n_splits)
        results[f'drop_{group_name}'] = metrics
        
        delta_auc = results['all_features']['auc_roc'] - metrics['auc_roc']
        print(f"  Drop {group_name:15s}: AUC={metrics['auc_roc']:.4f} "
              f"(Δ={delta_auc:+.4f}) | F1={metrics['f1_macro']:.4f}")
    
    # Only-one-group analysis
    print("\n--- Single Feature Group Performance ---")
    for group_name, group_cols in FEATURE_GROUPS.items():
        valid_cols = [c for c in group_cols if c in df.columns]
        if len(valid_cols) == 0:
            continue
        
        X_single = df[valid_cols].values
        metrics = _cross_validate(X_single, y, qt, n_splits)
        results[f'only_{group_name}'] = metrics
        print(f"  Only {group_name:15s}: AUC={metrics['auc_roc']:.4f} | F1={metrics['f1_macro']:.4f}")
    
    return results


def run_ablation_spatial_encoding(df, target_col='is_hotspot', n_splits=3):
    """Ablation 2: Impact of spatial encoding methods.
    
    Tests: no spatial, only spatial lag, only Fourier, only graph, only clusters, all spatial.
    """
    print("\n" + "="*80)
    print("ABLATION 2: Spatial Encoding Method Impact")
    print("="*80)
    
    y = df[target_col].values.astype(int)
    qt = QuantileTransformer(output_distribution='normal', random_state=42)
    results = {}
    
    configs = {
        'no_spatial': dict(include_spatial=False, include_fourier=False, 
                          include_graph=False, include_clusters=False),
        'spatial_lag_only': dict(include_spatial=True, include_fourier=False, 
                                include_graph=False, include_clusters=False),
        'fourier_only': dict(include_spatial=False, include_fourier=True, 
                            include_graph=False, include_clusters=False),
        'graph_only': dict(include_spatial=False, include_fourier=False, 
                          include_graph=True, include_clusters=False),
        'clusters_only': dict(include_spatial=False, include_fourier=False, 
                             include_graph=False, include_clusters=True),
        'all_spatial': dict(include_spatial=True, include_fourier=True, 
                           include_graph=True, include_clusters=True),
    }
    
    for config_name, config in configs.items():
        feature_cols = get_feature_columns(df, **config)
        X = df[feature_cols].replace([np.inf, -np.inf], np.nan).fillna(0).values
        
        metrics = _cross_validate(X, y, qt, n_splits)
        results[config_name] = metrics
        print(f"  {config_name:25s}: AUC={metrics['auc_roc']:.4f} | "
              f"F1={metrics['f1_macro']:.4f} | Features={len(feature_cols)}")
    
    return results


def run_ablation_clustering_granularity(df, target_col='is_hotspot', n_splits=3):
    """Ablation 3: Impact of clustering granularity (K=5,10,20,50)."""
    print("\n" + "="*80)
    print("ABLATION 3: Clustering Granularity Impact")
    print("="*80)
    
    y = df[target_col].values.astype(int)
    qt = QuantileTransformer(output_distribution='normal', random_state=42)
    results = {}
    
    base_features = get_feature_columns(df, include_spatial=True, include_fourier=True, 
                                        include_graph=True, include_clusters=False)
    
    for k in [5, 10, 20, 50]:
        cluster_cols = [c for c in df.columns if f'kmeans_{k}' in c]
        feature_cols = base_features + cluster_cols
        feature_cols = [f for f in feature_cols if f in df.columns]
        
        X = df[feature_cols].replace([np.inf, -np.inf], np.nan).fillna(0).values
        metrics = _cross_validate(X, y, qt, n_splits)
        results[f'k={k}'] = metrics
        print(f"  K={k:3d}: AUC={metrics['auc_roc']:.4f} | F1={metrics['f1_macro']:.4f}")
    
    # All cluster granularities combined
    cluster_cols = [c for c in df.columns if 'kmeans' in c or 'dbscan' in c or 'dist_to_center' in c]
    feature_cols = base_features + cluster_cols
    feature_cols = list(dict.fromkeys([f for f in feature_cols if f in df.columns]))
    X = df[feature_cols].replace([np.inf, -np.inf], np.nan).fillna(0).values
    metrics = _cross_validate(X, y, qt, n_splits)
    results['all_granularities'] = metrics
    print(f"  ALL : AUC={metrics['auc_roc']:.4f} | F1={metrics['f1_macro']:.4f}")
    
    return results


def run_ablation_transformer_embeddings(df, target_col='is_hotspot', n_splits=3):
    """Ablation 4: Transformer embeddings vs. pure GBDT."""
    print("\n" + "="*80)
    print("ABLATION 4: Transformer Embedding Impact")
    print("="*80)
    
    feature_cols = get_feature_columns(df)
    X = df[feature_cols].replace([np.inf, -np.inf], np.nan).fillna(0).values
    y = df[target_col].values.astype(int)
    
    results = {}
    
    # Pure GBDT (no transformer)
    print("\n  Training pure GBDT (no transformer)...")
    qt = QuantileTransformer(output_distribution='normal', random_state=42)
    metrics_gbdt = _cross_validate(X, y, qt, n_splits)
    results['pure_gbdt'] = metrics_gbdt
    print(f"  Pure GBDT: AUC={metrics_gbdt['auc_roc']:.4f} | F1={metrics_gbdt['f1_macro']:.4f}")
    
    # Hybrid with representative transformer configs (reduced for tractability)
    for d_token, n_layers in [(32, 2), (64, 2), (128, 2)]:
            config_name = f'hybrid_d{d_token}_L{n_layers}'
            print(f"\n  Training {config_name}...")
            
            metrics = _cross_validate_hybrid(X, y, n_splits, 
                                             d_token=d_token, n_layers=n_layers)
            results[config_name] = metrics
            print(f"  {config_name}: AUC={metrics['auc_roc']:.4f} | F1={metrics['f1_macro']:.4f}")
    
    return results


def run_ablation_calibration(df, target_col='is_hotspot', n_splits=3):
    """Ablation 5: Calibration method comparison."""
    print("\n" + "="*80)
    print("ABLATION 5: Calibration Method Comparison")
    print("="*80)
    
    feature_cols = get_feature_columns(df)
    X = df[feature_cols].replace([np.inf, -np.inf], np.nan).fillna(0).values
    y = df[target_col].values.astype(int)
    
    qt = QuantileTransformer(output_distribution='normal', random_state=42)
    
    # Get uncalibrated predictions via CV
    skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=42)
    uncal_probs = np.zeros(len(y))
    
    for fold, (train_idx, val_idx) in enumerate(skf.split(X, y)):
        X_train, X_val = X[train_idx], X[val_idx]
        y_train, y_val = y[train_idx], y[val_idx]
        
        X_train_qt = qt.fit_transform(X_train)
        X_val_qt = qt.transform(X_val)
        
        xgb_model = train_xgboost(X_train_qt, y_train, X_val_qt, y_val)
        lgb_model = train_lightgbm(X_train_qt, y_train, X_val_qt, y_val)
        
        xgb_probs = xgb_model.predict_proba(X_val_qt)[:, 1]
        lgb_probs = lgb_model.predict_proba(X_val_qt)[:, 1]
        uncal_probs[val_idx] = 0.5 * xgb_probs + 0.5 * lgb_probs
    
    # Now test each calibration method using nested CV
    results = {}
    
    # Uncalibrated
    ece_uncal, _ = expected_calibration_error(y, uncal_probs)
    results['uncalibrated'] = {
        'ECE': ece_uncal,
        'AUC': roc_auc_score(y, uncal_probs),
        'Brier': float(np.mean((uncal_probs - y)**2))
    }
    print(f"  Uncalibrated: ECE={ece_uncal:.4f} | AUC={results['uncalibrated']['AUC']:.4f}")
    
    # Test each calibration method
    for cal_name, CalClass in [('platt', PlattScaling), 
                                ('isotonic', IsotonicCalibration),
                                ('temperature', TemperatureScaling)]:
        cal_probs = np.zeros(len(y))
        
        skf_cal = StratifiedKFold(n_splits=3, shuffle=True, random_state=123)
        for _, (cal_train, cal_test) in enumerate(skf_cal.split(uncal_probs.reshape(-1, 1), y)):
            cal = CalClass()
            cal.fit(uncal_probs[cal_train], y[cal_train])
            cal_probs[cal_test] = cal.predict_proba(uncal_probs[cal_test])
        
        cal_probs = np.clip(cal_probs, 1e-7, 1-1e-7)
        ece_cal, _ = expected_calibration_error(y, cal_probs)
        results[cal_name] = {
            'ECE': ece_cal,
            'AUC': roc_auc_score(y, cal_probs),
            'Brier': float(np.mean((cal_probs - y)**2))
        }
        print(f"  {cal_name:12s}: ECE={ece_cal:.4f} | AUC={results[cal_name]['AUC']:.4f} | "
              f"Brier={results[cal_name]['Brier']:.4f}")
    
    return results


def _cross_validate(X, y, qt, n_splits=5):
    """Cross-validate with XGBoost+LightGBM ensemble."""
    from sklearn.metrics import roc_auc_score
    
    skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=42)
    all_metrics = []
    
    for fold, (train_idx, val_idx) in enumerate(skf.split(X, y)):
        X_train, X_val = X[train_idx], X[val_idx]
        y_train, y_val = y[train_idx], y[val_idx]
        
        qt_fold = QuantileTransformer(output_distribution='normal', random_state=42)
        X_train_qt = qt_fold.fit_transform(X_train)
        X_val_qt = qt_fold.transform(X_val)
        
        xgb_model = train_xgboost(X_train_qt, y_train, X_val_qt, y_val)
        lgb_model = train_lightgbm(X_train_qt, y_train, X_val_qt, y_val)
        
        xgb_probs = xgb_model.predict_proba(X_val_qt)[:, 1]
        lgb_probs = lgb_model.predict_proba(X_val_qt)[:, 1]
        probs = 0.5 * xgb_probs + 0.5 * lgb_probs
        preds = (probs >= 0.5).astype(int)
        
        fold_metrics = evaluate_model(y_val, preds, probs)
        all_metrics.append(fold_metrics)
    
    # Average metrics across folds
    avg_metrics = {}
    for key in all_metrics[0]:
        vals = [m[key] for m in all_metrics]
        avg_metrics[key] = float(np.mean(vals))
        avg_metrics[f'{key}_std'] = float(np.std(vals))
    
    return avg_metrics


def _cross_validate_hybrid(X, y, n_splits=5, d_token=64, n_layers=2):
    """Cross-validate with full hybrid model (transformer + GBDT)."""
    skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=42)
    all_metrics = []
    
    for fold, (train_idx, val_idx) in enumerate(skf.split(X, y)):
        X_train, X_val = X[train_idx], X[val_idx]
        y_train, y_val = y[train_idx], y[val_idx]
        
        model = HybridModel(
            task='classification', use_transformer=True, use_calibration=False,
            d_token=d_token, n_layers=n_layers
        )
        model.fit(X_train, y_train, X_val, y_val)
        
        probs = model.predict_proba(X_val)
        preds = (probs >= 0.5).astype(int)
        
        fold_metrics = evaluate_model(y_val, preds, probs)
        all_metrics.append(fold_metrics)
    
    avg_metrics = {}
    for key in all_metrics[0]:
        vals = [m[key] for m in all_metrics]
        avg_metrics[key] = float(np.mean(vals))
        avg_metrics[f'{key}_std'] = float(np.std(vals))
    
    return avg_metrics


def run_all_ablations(df, target_col='is_hotspot'):
    """Run all ablation studies and save results."""
    all_results = {}
    
    all_results['feature_groups'] = run_ablation_feature_groups(df, target_col)
    all_results['spatial_encoding'] = run_ablation_spatial_encoding(df, target_col)
    all_results['clustering_granularity'] = run_ablation_clustering_granularity(df, target_col)
    all_results['transformer_embeddings'] = run_ablation_transformer_embeddings(df, target_col)
    all_results['calibration'] = run_ablation_calibration(df, target_col)
    
    return all_results