#!/usr/bin/env python3
"""
Script to regenerate all standard plots (violin plots, ROC curves) for all experiments.
This ensures all experiments have the latest plot formatting with fixed x-axis labels.
"""

import argparse
from pathlib import Path
import glob
import os
from typing import Optional

from deception_detection.experiment import Experiment, ExperimentConfig
from deception_detection.postprocess import ExperimentPostprocessor
from deception_detection.prompt_scorer import PromptScorerMethod


def find_experiment_folders(search_dirs: Optional[list[str]] = None) -> list[Path]:
    """
    Find all experiment folders that contain cfg.yaml files.
    
    Args:
        search_dirs: List of directories to search in. Defaults to ['results', 'example_results']
    
    Returns:
        List of paths to experiment folders
    """
    if search_dirs is None:
        search_dirs = ['results', 'example_results']
    
    experiment_folders: list[Path] = []
    
    for search_dir in search_dirs:
        if not os.path.exists(search_dir):
            continue
            
        # Look for cfg.yaml files in subdirectories
        pattern = os.path.join(search_dir, '*', 'cfg.yaml')
        cfg_files = glob.glob(pattern)
        
        for cfg_file in cfg_files:
            exp_folder = Path(cfg_file).parent
            experiment_folders.append(exp_folder)
    
    return sorted(experiment_folders, key=lambda p: p.stat().st_mtime)


def regenerate_plots_for_experiment(exp_path: Path, scoring_method: PromptScorerMethod = "mean") -> bool:
    """
    Regenerate all standard plots for a single experiment.
    
    Args:
        exp_path: Path to the experiment folder
        scoring_method: Token reduction method to use
        
    Returns:
        True if successful, False otherwise
    """
    try:
        print(f"🔄 Regenerating plots for: {exp_path.name}")
        
        # Load experiment
        cfg = ExperimentConfig.from_path(exp_path)
        experiment = Experiment(cfg)
        
        # Run postprocessing to generate all plots
        postprocessor = ExperimentPostprocessor(
            experiment, 
            scoring_method=scoring_method,
            include_black_box=False
        )
        postprocessor.run()
        
        print(f"✅ Successfully regenerated plots for: {exp_path.name}")
        return True
        
    except Exception as e:
        print(f"❌ Error regenerating plots for {exp_path.name}: {e}")
        return False


def main():
    parser = argparse.ArgumentParser(description="Regenerate all standard plots for experiments")
    parser.add_argument("--search_dirs", type=str, nargs="+", default=["results"],
                       help="Directories to search for experiments (default: results)")
    parser.add_argument("--scoring_method", type=str, default="mean",
                       choices=["mean", "max", "relu", "topk", "final"],
                       help="Token reduction method to use (default: mean)")
    parser.add_argument("--list_experiments", action="store_true",
                       help="List all available experiment folders and exit")
    
    args = parser.parse_args()
    
    # List experiments if requested
    if args.list_experiments:
        experiment_folders = find_experiment_folders(args.search_dirs)
        print("Available experiment folders:")
        for i, folder in enumerate(experiment_folders, 1):
            print(f"  {i}. {folder}")
        return
    
    # Find all experiment folders
    experiment_folders = find_experiment_folders(args.search_dirs)
    
    if not experiment_folders:
        print("❌ No experiment folders found!")
        return
    
    print(f"🔍 Found {len(experiment_folders)} experiment folders")
    print(f"🎯 Using scoring method: {args.scoring_method}")
    print(f"📁 Searching in: {args.search_dirs}")
    print()
    
    # Regenerate plots for each experiment
    successful = 0
    failed = 0
    
    for i, exp_path in enumerate(experiment_folders, 1):
        print(f"\n{'='*60}")
        print(f"Processing experiment {i}/{len(experiment_folders)}: {exp_path.name}")
        print(f"{'='*60}")
        
        if regenerate_plots_for_experiment(exp_path, args.scoring_method):
            successful += 1
        else:
            failed += 1
    
    print(f"\n🎉 Plot regeneration complete!")
    print(f"✅ Successful: {successful}")
    print(f"❌ Failed: {failed}")
    print(f"📊 Total: {len(experiment_folders)}")
    
    if successful > 0:
        print(f"\n💡 Next steps:")
        print(f"   1. Run: python visualize_all_datasets.py --all_experiments --search_dirs results")
        print(f"   2. Run: python deception_detection/scripts/generate_html_dashboard.py --multi --search_dirs results")
        print(f"   3. Open the generated dashboard to see all visualizations!")


if __name__ == "__main__":
    main() 