#!/usr/bin/env python3
"""
Example showing how to train on multiple datasets by combining them manually.
"""

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

# Create your normal config but with one of the datasets
config = ExperimentConfig(
    method="lr",
    model_name=ModelName.LLAMA_70B_3_3,
    train_data="harm_refusal_choice__plain",  # This will be replaced
    eval_data=[
        "roleplaying__plain",
        "insider_trading__upscale", 
        "insider_trading_doubledown__upscale",
        "sandbagging_v2__wmdp_mmlu"
    ],
    id="multi_dataset_example",
    detect_layers=[20],
    reg_coeff=10
)

# Create the experiment
experiment = Experiment(config)

# Get the repository to load additional datasets
repository = DatasetRepository()

# Load multiple training datasets
dataset1 = repository.get(
    "harm_refusal_choice__plain", 
    model="prewritten",  # or use the model name if you want on-policy
    trim_reasoning=config.trim_reasoning
)

dataset2 = repository.get(
    "insider_trading__upscale",
    model="prewritten", 
    trim_reasoning=config.trim_reasoning
)

dataset3 = repository.get(
    "sandbagging_v2__wmdp_mmlu",
    model="prewritten",
    trim_reasoning=config.trim_reasoning
)

# Combine the datasets using the + operator
combined_training_dataset = dataset1 + dataset2 + dataset3

# Apply validation split manually
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

# Replace the experiment's training dataset
experiment.train_dataset = train_split

print(f"Combined training dataset size: {len(experiment.train_dataset)}")
print(f"Dataset names: {combined_training_dataset.base_name}")
print(f"Dataset variant: {combined_training_dataset.variant}")

# Now you can train and evaluate as normal
detector = experiment.get_detector()
experiment.compute_eval_scores()

# Your experiment will now be trained on all three datasets combined! 