#!/usr/bin/env python3
"""
Configuration Matrix Generator for Multi-Agent Maze Runner

This script generates configuration files by creating a matrix exploration of different
maze IDs and temperature values. Modify the lists below to change which parameters
are explored in the generated configurations.

Usage:
    python generate_configs.py

Generated files will be saved to the 'generated_configs/' directory.
"""

import os
import yaml
import random
from pathlib import Path
from typing import List, Dict, Any

# ================== CONFIGURABLE PARAMETERS ==================

# Maze IDs by difficulty level (extracted from collection_metadata.json)
EASY_MAZES: List[str] = [
    "a4b8a6712204421e",  # Easy maze
    "2098199777794f7e",  # Easy maze
    "849bb4f93b784776",  # Easy maze
    "412eed728b314ae1",  # Easy maze
    "2f94ec7921074e51",  # Easy maze
]

MEDIUM_MAZES: List[str] = [
    "cb72cc3b86684737",  # Medium maze
    "d1d68192f1394d4d",  # Medium maze
    "8e549088ecfa455c",  # Medium maze
    "d82959a5f0f64dcc",  # Medium maze
    "2ea1d0fbcf2c49b4",  # Medium maze
]

HARD_MAZES: List[str] = [
    "fc07a116411640d3",  # Hard maze
    "83233df81812411c",  # Hard maze
    "31a87865d40c4f65",  # Hard maze
    "3b16d8e6011d49ab",  # Hard maze
    "a7543add40404891",  # Hard maze
]

VERY_HARD_MAZES: List[str] = [
    "4c75bae517da44a2",  # Very hard maze
    "c41c16219cee401a",  # Very hard maze
    "5eefafe4b7bf46da",  # Very hard maze
    "c795e060f629476d",  # Very hard maze
    "fc39d5b12a8c4e69",  # Very hard maze
]

# Experiment configurations
EXPERIMENT_CONFIGS = [
    # BASELINE CONFIG: 1 agent, gpt-5 models
    {
        "name": "baseline",
        "num_execution_agents": 1,
        "execution_model": "gpt-5",
        "planning_model": "gpt-4.1-nano", 
        "orchestration_model": "gpt-5",
    },
    # N-AGENTS CONFIG A: 2 agents, gpt-4.1-nano models
    {
        "name": "n_agents_2_nano", 
        "num_execution_agents": 2,
        "execution_model": "gpt-4.1-nano",
        "planning_model": "gpt-4.1-nano",
        "orchestration_model": "gpt-4.1-nano",
    },
    # N-AGENTS CONFIG B: 3 agents, gpt-4.1-nano models
    {
        "name": "n_agents_3_nano",
        "num_execution_agents": 3, 
        "execution_model": "gpt-4.1-nano",
        "planning_model": "gpt-4.1-nano",
        "orchestration_model": "gpt-4.1-nano",
    },
    # N-AGENTS CONFIG C: 2 agents, mixed models (nano execution, gpt-5 orchestration)
    {
        "name": "n_agents_2_mixed",
        "num_execution_agents": 2,
        "execution_model": "gpt-4.1-nano", 
        "planning_model": "gpt-4.1-nano",
        "orchestration_model": "gpt-5",
    },
]

# Ablation configurations - 3 different ablation settings per config
ABLATION_CONFIGS = [
    {
        "name": "no_ablations",
        "enable_orchestration_agent": False,
        "enable_dynamic_weights": False, 
        "enable_teammate_coordination": False,
        "enable_reward_function": False,
    },
    {
        "name": "weights_and_reward",
        "enable_orchestration_agent": False,
        "enable_dynamic_weights": True,
        "enable_teammate_coordination": False, 
        "enable_reward_function": True,
    },
    {
        "name": "full_ablations",
        "enable_orchestration_agent": True,
        "enable_dynamic_weights": True,
        "enable_teammate_coordination": True,
        "enable_reward_function": True,
    },
]

# Base configuration file to use as template
BASE_CONFIG_PATH = "config.yaml"

# Output directory for generated configurations
OUTPUT_DIR = "generated_configs"

# Number of runs per configuration (for repeated experiments)
RUNS_PER_CONFIG = 3

# ================== SCRIPT IMPLEMENTATION ==================

def select_test_mazes() -> Dict[str, str]:
    """
    Randomly select one medium and one hard maze for experiments.
    
    Returns:
        Dictionary with selected maze IDs for each difficulty
    """
    # Set random seed for reproducible selection (can be changed for different selections)
    random.seed(42)  # Change this number for different maze selections
    
    selected_mazes = {
        "medium": random.choice(MEDIUM_MAZES),
        "hard": random.choice(HARD_MAZES)
    }
    
    print("Selected mazes for experiments:")
    print(f"  Medium difficulty: {selected_mazes['medium']}")  
    print(f"  Hard difficulty: {selected_mazes['hard']}")
    
    return selected_mazes

def load_base_config(config_path: str) -> Dict[str, Any]:
    """
    Load the base configuration file as a template.
    
    Args:
        config_path: Path to the base config.yaml file
        
    Returns:
        Dictionary containing the base configuration
    """
    with open(config_path, 'r') as f:
        return yaml.safe_load(f)

def generate_config_filename(maze_id: str, experiment_config: Dict[str, Any], 
                           ablation_config: Dict[str, Any], run_number: int) -> str:
    """
    Generate a descriptive filename for the configuration.
    
    Args:
        maze_id: UUID of the maze
        experiment_config: Experiment configuration dict
        ablation_config: Ablation configuration dict
        run_number: Run number (1, 2, 3)
        
    Returns:
        Formatted filename string
    """
    exp_name = experiment_config["name"]
    ablation_name = ablation_config["name"]
    return f"config_{maze_id}_{exp_name}_{ablation_name}_run{run_number}.yaml"

def create_config_variant(base_config: Dict[str, Any], maze_id: str, 
                         experiment_config: Dict[str, Any], 
                         ablation_config: Dict[str, Any], run_number: int) -> Dict[str, Any]:
    """
    Create a configuration variant with specified settings.
    
    Args:
        base_config: Base configuration dictionary
        maze_id: UUID of the maze to use
        experiment_config: Experiment configuration
        ablation_config: Ablation configuration
        run_number: Run number for this configuration
        
    Returns:
        Modified configuration dictionary
    """
    # Deep copy the base config to avoid modifying the original
    import copy
    config = copy.deepcopy(base_config)
    
    # Update maze ID
    config['maze']['uuid'] = maze_id
    
    # Update agent configurations
    config['agents']['global']['num_execution_agents'] = experiment_config['num_execution_agents']
    config['agents']['execution']['model'] = experiment_config['execution_model']
    config['agents']['planning']['model'] = experiment_config['planning_model']
    config['agents']['orchestration']['model'] = experiment_config['orchestration_model']
    
    # Update ablation settings
    config['ablations']['enable_orchestration_agent'] = ablation_config['enable_orchestration_agent']
    config['ablations']['enable_dynamic_weights'] = ablation_config['enable_dynamic_weights']
    config['ablations']['enable_teammate_coordination'] = ablation_config['enable_teammate_coordination']
    config['ablations']['enable_reward_function'] = ablation_config['enable_reward_function']
    
    # Disable visualization by default for batch experiments
    if 'visualization' in config:
        config['visualization']['enable_maze_visualization'] = False
    
    # Keep temperature from base config (0.1)
    # Keep other settings from base config unchanged
    
    return config

def save_config(config: Dict[str, Any], filepath: str, maze_id: str, 
               experiment_config: Dict[str, Any], ablation_config: Dict[str, Any], 
               run_number: int) -> None:
    """
    Save configuration to a YAML file with proper formatting.
    
    Args:
        config: Configuration dictionary to save
        filepath: Path where to save the file
        maze_id: UUID of the maze
        experiment_config: Experiment configuration
        ablation_config: Ablation configuration
        run_number: Run number
    """
    with open(filepath, 'w') as f:
        # Add header comment
        f.write("# Generated configuration file for ablation study\n")
        f.write(f"# Maze ID: {maze_id}\n")
        f.write(f"# Experiment: {experiment_config['name']}\n")
        f.write(f"# Agents: {experiment_config['num_execution_agents']}\n")
        f.write(f"# Models: {experiment_config['execution_model']} / {experiment_config['orchestration_model']}\n")
        f.write(f"# Ablation: {ablation_config['name']}\n") 
        f.write(f"# Run: {run_number}/{RUNS_PER_CONFIG}\n")
        f.write("# Generated by: generate_configs.py\n\n")
        
        # Write the YAML content
        yaml.dump(config, f, default_flow_style=False, indent=2, sort_keys=False)

def generate_all_configs() -> None:
    """
    Generate all configuration combinations for the ablation study.
    """
    # Create output directory if it doesn't exist
    Path(OUTPUT_DIR).mkdir(exist_ok=True)
    
    # Load base configuration
    try:
        base_config = load_base_config(BASE_CONFIG_PATH)
        print(f"✓ Loaded base configuration from: {BASE_CONFIG_PATH}")
    except FileNotFoundError:
        print(f"✗ Error: Base configuration file not found: {BASE_CONFIG_PATH}")
        return
    except yaml.YAMLError as e:
        print(f"✗ Error: Failed to parse base configuration: {e}")
        return
    
    # Select test mazes
    selected_mazes = select_test_mazes()
    
    # Calculate total combinations
    total_combinations = len(selected_mazes) * len(EXPERIMENT_CONFIGS) * len(ABLATION_CONFIGS) * RUNS_PER_CONFIG
    print(f"Generating {total_combinations} configuration files...")
    print(f"Mazes: {len(selected_mazes)} (medium + hard)")
    print(f"Experiment configs: {len(EXPERIMENT_CONFIGS)}")
    print(f"Ablation configs: {len(ABLATION_CONFIGS)}")
    print(f"Runs per config: {RUNS_PER_CONFIG}")
    
    generated_count = 0
    
    for difficulty, maze_id in selected_mazes.items():
        for experiment_config in EXPERIMENT_CONFIGS:
            for ablation_config in ABLATION_CONFIGS:
                for run_number in range(1, RUNS_PER_CONFIG + 1):
                    try:
                        # Create configuration variant
                        config = create_config_variant(base_config, maze_id, 
                                                     experiment_config, ablation_config, run_number)
                        
                        # Generate filename and path
                        filename = generate_config_filename(maze_id, experiment_config, 
                                                          ablation_config, run_number)
                        filepath = os.path.join(OUTPUT_DIR, filename)
                        
                        # Save configuration
                        save_config(config, filepath, maze_id, experiment_config, 
                                  ablation_config, run_number)
                        generated_count += 1
                        
                        if generated_count % 10 == 0:  # Progress indicator
                            print(f"Generated {generated_count}/{total_combinations} configurations...")
                            
                    except Exception as e:
                        print(f"✗ Error generating config for {maze_id}, {experiment_config['name']}, "
                              f"{ablation_config['name']}, run {run_number}: {e}")
    
    print(f"✓ Successfully generated {generated_count} configuration files in '{OUTPUT_DIR}/' directory")
    
    # Print detailed summary
    print("\nExperiment Summary:")
    print("  Selected mazes:")
    for difficulty, maze_id in selected_mazes.items():
        print(f"    {difficulty.title()}: {maze_id}")
    print(f"  Experiment configurations: {len(EXPERIMENT_CONFIGS)}")
    for exp_config in EXPERIMENT_CONFIGS:
        print(f"    - {exp_config['name']}: {exp_config['num_execution_agents']} agents, "
              f"{exp_config['execution_model']}/{exp_config['orchestration_model']}")
    print(f"  Ablation settings: {len(ABLATION_CONFIGS)}")
    for abl_config in ABLATION_CONFIGS:
        print(f"    - {abl_config['name']}: {abl_config}")
    print(f"  Runs per configuration: {RUNS_PER_CONFIG}")
    print(f"  Total configurations: {generated_count}")
    print(f"  Output directory: {OUTPUT_DIR}/")

def print_experiment_design() -> None:
    """
    Print detailed information about the experimental design.
    """
    print("\nExperimental Design:")
    print("=" * 60)
    print("BASELINE CONFIG: 1 execution agent")
    print("  Model: gpt-5 (execution & orchestration)")
    print("  Planning: gpt-4.1-nano")
    print("  - 3 runs with all ablations false")  
    print("  - 3 runs with dynamic_weights + reward_function enabled")
    print("  - 3 runs with dynamic_weights + reward_function + orchestration enabled")
    print()
    
    print("N-AGENTS CONFIG A: 2 execution agents")
    print("  Models: gpt-4.1-nano (all)")
    print("  - 3 runs with all ablations false")
    print("  - 3 runs with dynamic_weights + reward_function enabled") 
    print("  - 3 runs with dynamic_weights + reward_function + orchestration enabled")
    print()
    
    print("N-AGENTS CONFIG B: 3 execution agents")
    print("  Models: gpt-4.1-nano (all)")
    print("  - 3 runs with all ablations false")
    print("  - 3 runs with dynamic_weights + reward_function enabled")
    print("  - 3 runs with dynamic_weights + reward_function + orchestration enabled")
    print()
    
    print("N-AGENTS CONFIG C: 2 execution agents")  
    print("  Models: gpt-4.1-nano (execution & planning), gpt-5 (orchestration)")
    print("  - 3 runs with all ablations false")
    print("  - 3 runs with dynamic_weights + reward_function enabled")
    print("  - 3 runs with dynamic_weights + reward_function + orchestration enabled")
    print()
    
    print("Each configuration tested on:")
    print("  - 1 randomly selected medium difficulty maze")
    print("  - 1 randomly selected hard difficulty maze")
    print(f"Total configurations: 4 experiment configs × 3 ablation configs × 3 runs × 2 mazes = {4*3*3*2} configs")

if __name__ == "__main__":
    print("Multi-Agent Maze Runner - Ablation Study Configuration Generator")
    print("=" * 80)
    
    # Print experimental design
    print_experiment_design()
    
    # Generate all configurations
    print("\nStarting configuration generation...")
    generate_all_configs()
    
    print("\nDone! You can now run batch experiments using:")
    print(f"  python batch_runner.py --config_dir {OUTPUT_DIR} --parallel 4")
    print("\nTo run individual configs:")
    print(f"  python main.py --config {OUTPUT_DIR}/config_[maze_id]_[exp_name]_[ablation]_run1.yaml")