#!/usr/bin/env python3
"""
Update existing results_table.csv files with balanced accuracy at threshold=0.
This script adds a new column 'balanced_accuracy_at_0' to each experiment's results_table.csv.
For single-class datasets, it will use regular accuracy instead of balanced accuracy.
"""

import os
import pandas as pd
import torch
from pathlib import Path

# Set environment variables to reduce memory usage
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["CUDA_VISIBLE_DEVICES"] = "1,7,6"
from deception_detection.experiment import Experiment, ExperimentConfig
from deception_detection.prompt_scorer import create_prompt_scorer


def update_results_table(experiment_folder: Path) -> bool:
    """Update the results_table.csv file with balanced accuracy at threshold=0."""
    # Check if results_table.csv exists
    results_table_path = experiment_folder / "results_table.csv"
    if not results_table_path.exists():
        print("  ⚠️  No results_table.csv found")
        return False
    
    # Clear GPU memory before loading experiment
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    
    # Load experiment
    cfg = ExperimentConfig.from_path(experiment_folder)
    experiment = Experiment(cfg)
    type = "max"
    # Create scorer WITH experiment so control dataset is loaded for FPR calculations
    scorer = create_prompt_scorer(
        method=type, 
        experiment=experiment, 
        fprs=[0.01, 0.001], 
        default_fpr=0.01, 
        control_dataset_name="alpaca_XYZ__plain"
    )
    
    # Get all evaluation scores
    eval_scores = experiment.get_eval_scores()

    # Load existing results table
    results_df = pd.read_csv(results_table_path, index_col=0)
    
    # Check if column already exists
    if f'balanced_accuracy_at_0_{type}' in results_df.columns:
        print("  ℹ️  Column 'balanced_accuracy_at_0' already exists, updating values...")
    
    # Calculate balanced accuracy at threshold=0 for each dataset
    ba_at_zero_values = []
    ba_at_fpr_values = []
    for dataset_name in results_df.index:
        if dataset_name in eval_scores:
            print("######## DATASET NAME: ", dataset_name)
            # Use the updated balanced_accuracy method (handles single-class datasets)
            ba_at_zero = scorer.balanced_accuracy(eval_scores[dataset_name])
            ba_at_zero_values.append(ba_at_zero)
            print(f"######    {dataset_name} accuracy: {ba_at_zero:.4f}")

            balanced_accuracy_at_fpr = scorer.balanced_accuracy_at_fpr(eval_scores[dataset_name])
            ba_at_fpr_values.append(balanced_accuracy_at_fpr)
            print(f"######    {dataset_name} balanced_accuracy_at_fpr: {balanced_accuracy_at_fpr:.4f}")

        else:
            ba_at_zero_values.append(float("nan"))
            print(f"######    {dataset_name}: NaN (no scores found)")
    
    # Add/update column in results table
    results_df[f'balanced_accuracy_at_0_{type}'] = ba_at_zero_values
    results_df[f'balanced_accuracy_at_fpr_{type}'] = ba_at_fpr_values
    # Save updated results table
    results_df.to_csv(results_table_path)
    print("  ✅ Updated results_table.csv")
    
    # Clean up memory after processing
    del experiment, scorer, eval_scores, results_df
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    
    return True



def main():
    """Main function to update all results_table.csv files."""
    results_dir = Path("results")
    
    if not results_dir.exists():
        print(f"Results directory {results_dir} not found!")
        return
    
    experiment_folders = [f for f in results_dir.iterdir() if f.is_dir()]
    
    print(f"🔍 Found {len(experiment_folders)} experiment folders")
    print(f"📝 Updating results_table.csv files with balanced_accuracy_at_0_{type} column...")
    print("=" * 80)
    
    updated_count = 0
    skipped_count = 0
    
    for experiment_folder in experiment_folders:
        print(f"\n📁 {experiment_folder.name}")
        
        if update_results_table(experiment_folder):
            updated_count += 1
        else:
            skipped_count += 1
            
        # Clear GPU memory after processing each experiment
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    
    print("\n" + "=" * 80)
    print("📊 Summary:")
    print(f"  ✅ Updated: {updated_count} experiments")
    print(f"  ⏭️  Skipped: {skipped_count} experiments")
    print(f"  📈 Total processed: {updated_count + skipped_count} experiments")


if __name__ == "__main__":
    main()
