#!/usr/bin/env python3
"""
Script to train a detector on multiple combined datasets.
Based on the upper_bound.yaml config but with multiple training datasets.
"""

from deception_detection.experiment import Experiment, ExperimentConfig
from deception_detection.repository import DatasetRepository
from deception_detection.models import ModelName

def create_multi_dataset_experiment():
    """Create an experiment that trains on multiple combined datasets."""
    
    # Base configuration (based on XXXX-1.yaml but with multiple training datasets)
    config = ExperimentConfig(
        method="lr",
        model_name=ModelName.MISTRAL_3_1_24B,
        train_data="repe_honesty__plain",  # This will be replaced with combined dataset
        eval_data=[
            "harm_refusal_knowledge_report__plain",
            "instruct_dishonesty_default__plain",
            "instruct_dishonesty_donotadmit__plain",
            "mask__plain",
            "insider_trading_XYZ__plain"
        ],
        eval_on_policy=False,
        control_data=["alpaca_XYZ__plain"],
        id="multi_dataset_XYZ",
        detect_only_start_of_turn=False,
        val_fraction=0.2,
        control_on_policy=False,
        control_dataset_size=2000,
        detect_layers=[20],
        reg_coeff=10
    )
    
    # Create the experiment
    experiment = Experiment(config)
    
    # Get the repository to load additional datasets
    repository = DatasetRepository()
    
    # List of datasets you want to train on (from XXXX-1.yaml + the commented one for mistral)
    training_dataset_names = [
        # "harm_refusal_choice__plain",  # The commented one that works only for mistral
        "harm_refusal_knowledge_report__plain",
        "instruct_dishonesty_default__plain",
        "instruct_dishonesty_donotadmit__plain", 
        "mask__plain",
        "insider_trading_XYZ__plain"
    ]
    
    print(f"Loading {len(training_dataset_names)} training datasets for mistral...")
    
    # Load all training datasets
    training_datasets = []
    for dataset_name in training_dataset_names:
        print(f"  Loading {dataset_name}...")
        dataset = repository.get(
            dataset_name,
            model="prewritten",  # Change to experiment.train_policy_model if you want on-policy
            trim_reasoning=config.trim_reasoning,
            shuffle_upon_init=False,  # We'll shuffle the combined dataset
            max_llama_token_length=config.max_llama_token_length
        )
        training_datasets.append(dataset)
        print(f"    Size: {len(dataset)} dialogues")
    
    # Combine all training datasets
    print("\nCombining datasets...")
    combined_training_dataset = training_datasets[0]
    for dataset in training_datasets[1:]:
        combined_training_dataset = combined_training_dataset + dataset
    
    print(f"Combined dataset size: {len(combined_training_dataset)} dialogues")
    print(f"Combined dataset name: {combined_training_dataset.base_name}")
    
    # Apply validation split manually (same logic as get_train_dataset)
    num_val = int(len(combined_training_dataset) * config.val_fraction)
    train_split = combined_training_dataset[:-num_val] if num_val > 0 else combined_training_dataset
    
    print(f"Training split size: {len(train_split)} dialogues")
    if num_val > 0:
        print(f"Validation split size: {num_val} dialogues")
    
    # Replace the experiment's training dataset
    experiment.train_dataset = train_split
    
    return experiment

def main():
    """Main function to run the multi-dataset training."""
    print("=== Multi-Dataset Training Setup ===")
    
    # Create the experiment with combined training data
    experiment = create_multi_dataset_experiment()
    
    print(f"\n=== Training detector on combined dataset ===")
    print(f"Method: {experiment.cfg.method}")
    print(f"Model: {experiment.cfg.model_name.value}")
    print(f"Training on: {len(experiment.train_dataset)} dialogues")
    print(f"Reg coefficient: {experiment.cfg.reg_coeff}")
    
    # Train the detector
    _ = experiment.get_detector()
    print("✓ Detector trained successfully!")
    
    print(f"\n=== Computing evaluation scores ===")
    # Compute evaluation scores
    experiment.compute_eval_scores()
    print("✓ Evaluation scores computed!")
    
    # Get results
    from deception_detection.postprocess import ExperimentPostprocessor
    postprocessor = ExperimentPostprocessor(experiment)
    results_df = postprocessor.get_results_dataframe()
    
    print(f"\n=== Results ===")
    print(results_df[['auroc', 'balanced_accuracy']].round(3))
    
    print(f"\nResults saved to: {experiment.cfg.folder}")
    print(f"Experiment name: {experiment.cfg.name}")

if __name__ == "__main__":
    main() 