#!/usr/bin/env python3
"""
Script to visualize probe predictions for all evaluation datasets in an experiment.

Usage:
    python visualize_all_datasets.py [--exp_path /path/to/experiment/folder] [--output_dir /path/to/output]
"""

import argparse
from pathlib import Path
import torch as t
import matplotlib.pyplot as plt
import glob
import os
from typing import Optional, Any

from deception_detection.experiment import Experiment, ExperimentConfig
from deception_detection.prompt_scorer import create_prompt_scorer
from deception_detection.types import Label


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 get_latest_experiment_folder(search_dirs: Optional[list[str]] = None) -> Path:
    """
    Get the most recently created experiment folder.
    
    Args:
        search_dirs: List of directories to search in
        
    Returns:
        Path to the latest experiment folder
    """
    experiment_folders = find_experiment_folders(search_dirs)
    
    if not experiment_folders:
        raise ValueError("No experiment folders found with cfg.yaml files")
    
    # Sort by modification time (most recent first)
    experiment_folders.sort(key=lambda p: p.stat().st_mtime, reverse=True)
    
    return experiment_folders[0]


def visualize_preds(
    scores: t.Tensor, 
    labels: t.Tensor, 
    output_dir: Path, 
    model_name: str, 
    layer: str, 
    token: str, 
    dataset_name: str
) -> None:
    """
    Visualize predictions vs labels using a scatter plot.
    
    Args:
        scores: Tensor of probe scores
        labels: Tensor of true labels (boolean, True = deceptive)
        output_dir: Directory to save the plot
        model_name: Name of the model
        layer: Layer number(s)
        token: Token type (e.g., "detection_mask")
        dataset_name: Name of the dataset
    """
    # Convert tensors to numpy arrays
    scores_np = scores.to(t.float32).cpu().numpy()
    mask_np = labels.to(t.float32).bool().cpu().numpy()

    # Create a new figure
    plt.figure(figsize=(10, 6))

    # Cast to Any to avoid type checker warnings about numpy arrays
    scores_honest: Any = scores_np[~mask_np]
    scores_dishonest: Any = scores_np[mask_np]

    plt.hist(scores_dishonest, bins=25, label="deceptive", alpha=0.5, color="red")
    plt.hist(scores_honest, bins=25, label="honest", alpha=0.5, color="blue")

    plt.title(f"Probe scores for {model_name} layer {layer} token {token}\nDataset: {dataset_name}")
    plt.xlabel("Probe Score")
    plt.ylabel("Frequency")
    plt.legend()

    # save image
    output_dir.mkdir(parents=True, exist_ok=True)
    path = Path(output_dir) / f"probe_scores_{model_name}_layer_{layer}_{token}_{dataset_name}.png".replace("/", "_")
    plt.savefig(path, dpi=300, bbox_inches='tight')
    plt.close()  # Close the figure to free memory and prevent legend accumulation
    print(f"Saved plot: {path}")


def visualize_all_datasets(exp_path: str, output_dir: Optional[str] = None, token_types: Optional[list[str]] = None):
    """
    Visualize probe predictions for all evaluation datasets in an experiment.
    
    Args:
        exp_path: Path to the experiment folder
        output_dir: Directory to save the plots (if None, saves in experiment folder)
        token_types: List of token types to use (e.g., ["mean", "max"])
    """
    if token_types is None:
        token_types = ["mean", "max"]
    
    # Load experiment
    cfg = ExperimentConfig.from_path(Path(exp_path))
    experiment = Experiment(cfg)
    
    # Get model name and layers
    model_name = experiment.cfg.model_name.value
    layers_str = "_".join(map(str, experiment.cfg.detect_layers))
    
    print(f"Processing experiment: {exp_path}")
    print(f"Model: {model_name}")
    print(f"Layers: {layers_str}")
    
    # Get scores for processing
    all_scores = experiment.get_eval_scores(all_tokens=False)
    
    # Set default output directory to experiment folder if not provided
    if output_dir is None:
        output_dir = str(Path(exp_path) / "plots")
    
    print(f"Saving plots to: {output_dir}")
    
    # Process each token type
    for token_type in token_types:
        print(f"\n=== Processing token type: {token_type} ===")
        
        # Create prompt scorer based on token type
        scorer = create_prompt_scorer(token_type, experiment, fprs=[0.01], control_dataset_name="alpaca_XYZ__plain")  # type: ignore
        
        print(f"Found {len(all_scores)} evaluation datasets:")
        for dataset_name in all_scores.keys():
            print(f"  - {dataset_name}")
        
        # Process each evaluation dataset
        for dataset_name, dataset_scores in all_scores.items():
            print(f"\nProcessing dataset: {dataset_name}")
            
            try:
                # Get the reduced scores (one score per dialogue) using the specified method
                reduced_scores = scorer.process(dataset_scores, normalize=False)
                
                # Convert labels to boolean tensor (True = deceptive)
                label_tensor = t.tensor([label == Label.DECEPTIVE for label in dataset_scores.labels])
                
                # Create visualization
                visualize_preds(
                    scores=reduced_scores,
                    labels=label_tensor,
                    output_dir=Path(output_dir) / token_type,
                    model_name=model_name,
                    layer=layers_str,
                    token=token_type,
                    dataset_name=dataset_name
                )
                
            except Exception as e:
                print(f"Error processing dataset {dataset_name}: {e}")
                continue
    
    print(f"\nAll plots saved to: {output_dir}")


def main():
    parser = argparse.ArgumentParser(description="Visualize probe predictions for all evaluation datasets")
    parser.add_argument("--exp_path", type=str, default=None,
                       help="Path to the experiment folder (if not provided, will use the latest experiment)")
    parser.add_argument("--output_dir", type=str, default=None,
                       help="Directory to save the visualization plots (if not provided, saves in experiment folder)")
    parser.add_argument("--token_types", type=str, nargs="+", default=["mean", "max"],
                       choices=["mean", "max", "relu", "topk", "final"],
                       help="Types of token reduction methods to use")
    parser.add_argument("--all_experiments", action="store_true",
                       help="Process all experiment folders instead of just one")
    parser.add_argument("--search_dirs", type=str, nargs="+", default=["results", "example_results"],
                       help="Directories to search for experiments (default: results example_results)")
    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
    
    # Process all experiments if requested
    if args.all_experiments:
        experiment_folders = find_experiment_folders(args.search_dirs)
        print(f"Processing {len(experiment_folders)} experiments...")
        
        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}")
            
            try:
                visualize_all_datasets(
                    exp_path=str(exp_path),
                    output_dir=args.output_dir,
                    token_types=args.token_types
                )
                print(f"✅ Completed: {exp_path.name}")
            except Exception as e:
                print(f"❌ Error processing {exp_path.name}: {e}")
                continue
        
        print(f"\n🎉 Finished processing all {len(experiment_folders)} experiments!")
        return
    
    # Process single experiment (existing behavior)
    if args.exp_path is None:
        try:
            exp_path = get_latest_experiment_folder(args.search_dirs)
            print(f"Auto-detected latest experiment: {exp_path}")
        except ValueError as e:
            print(f"Error: {e}")
            print("Use --list_experiments to see available experiments")
            return
    else:
        exp_path = Path(args.exp_path)
        if not exp_path.exists():
            raise ValueError(f"Experiment path does not exist: {exp_path}")
        
        config_file = exp_path / "cfg.yaml"
        if not config_file.exists():
            raise ValueError(f"Config file not found: {config_file}")
    
    # Create output directory if specified and it doesn't exist
    if args.output_dir is not None:
        output_dir = Path(args.output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
    
    # Run visualization
    visualize_all_datasets(
        exp_path=str(exp_path),
        output_dir=args.output_dir,
        token_types=args.token_types
    )


if __name__ == "__main__":
    main() 