DynamicAbstractionSystem/config/config_loader.py

"""Configuration loading and management."""

import json
import yaml
from pathlib import Path
from typing import Dict, Any, Union
from dataclasses import asdict

from .simulation_config import HeadlessConfig, InteractiveConfig, ExperimentConfig


class ConfigLoader:
    """Loads configuration from files and creates config objects."""

    @staticmethod
    def load_headless_config(config_path: Union[str, Path]) -> HeadlessConfig:
        """Load headless configuration from file."""
        config_path = Path(config_path)

        if not config_path.exists():
            raise FileNotFoundError(f"Config file not found: {config_path}")

        with open(config_path, 'r') as f:
            if config_path.suffix.lower() == '.json':
                data = json.load(f)
            elif config_path.suffix.lower() in ['.yml', '.yaml']:
                data = yaml.safe_load(f)
            else:
                raise ValueError(f"Unsupported config format: {config_path.suffix}")

        return ConfigLoader._dict_to_headless_config(data)

    @staticmethod
    def load_interactive_config(config_path: Union[str, Path]) -> InteractiveConfig:
        """Load interactive configuration from file."""
        config_path = Path(config_path)

        if not config_path.exists():
            # Return default config if file doesn't exist
            return InteractiveConfig()

        with open(config_path, 'r') as f:
            if config_path.suffix.lower() == '.json':
                data = json.load(f)
            elif config_path.suffix.lower() in ['.yml', '.yaml']:
                data = yaml.safe_load(f)
            else:
                raise ValueError(f"Unsupported config format: {config_path.suffix}")

        return ConfigLoader._dict_to_interactive_config(data)

    @staticmethod
    def load_experiment_config(config_path: Union[str, Path]) -> ExperimentConfig:
        """Load experiment configuration from file."""
        config_path = Path(config_path)

        if not config_path.exists():
            raise FileNotFoundError(f"Config file not found: {config_path}")

        with open(config_path, 'r') as f:
            if config_path.suffix.lower() == '.json':
                data = json.load(f)
            elif config_path.suffix.lower() in ['.yml', '.yaml']:
                data = yaml.safe_load(f)
            else:
                raise ValueError(f"Unsupported config format: {config_path.suffix}")

        return ConfigLoader._dict_to_experiment_config(data)

    @staticmethod
    def save_config(config: Union[HeadlessConfig, InteractiveConfig, ExperimentConfig],
                    config_path: Union[str, Path]):
        """Save configuration to file."""
        config_path = Path(config_path)
        config_path.parent.mkdir(parents=True, exist_ok=True)

        # Convert config to dictionary
        if isinstance(config, HeadlessConfig):
            data = asdict(config)
        elif isinstance(config, InteractiveConfig):
            data = asdict(config)
        elif isinstance(config, ExperimentConfig):
            data = asdict(config)
        else:
            raise ValueError(f"Unknown config type: {type(config)}")

        with open(config_path, 'w') as f:
            if config_path.suffix.lower() == '.json':
                json.dump(data, f, indent=2)
            elif config_path.suffix.lower() in ['.yml', '.yaml']:
                yaml.dump(data, f, default_flow_style=False)
            else:
                raise ValueError(f"Unsupported config format: {config_path.suffix}")

    @staticmethod
    def _dict_to_headless_config(data: Dict[str, Any]) -> HeadlessConfig:
        """Convert dictionary to HeadlessConfig."""
        # Extract output config if present
        output_data = data.get('output', {})
        output_config = OutputConfig(**output_data)

        # Extract simulation config if present
        sim_data = data.get('simulation', {})
        simulation_config = SimulationConfig(**sim_data)

        return HeadlessConfig(
            max_ticks=data.get('max_ticks'),
            max_duration=data.get('max_duration'),
            output=output_config,
            simulation=simulation_config
        )

    @staticmethod
    def _dict_to_interactive_config(data: Dict[str, Any]) -> InteractiveConfig:
        """Convert dictionary to InteractiveConfig."""
        # Extract simulation config if present
        sim_data = data.get('simulation', {})
        simulation_config = SimulationConfig(**sim_data)

        return InteractiveConfig(
            window_width=data.get('window_width', 0),
            window_height=data.get('window_height', 0),
            vsync=data.get('vsync', True),
            resizable=data.get('resizable', True),
            show_grid=data.get('show_grid', True),
            show_interaction_radius=data.get('show_interaction_radius', False),
            show_legend=data.get('show_legend', True),
            control_bar_height=data.get('control_bar_height', 48),
            inspector_width=data.get('inspector_width', 260),
            properties_width=data.get('properties_width', 320),
            console_height=data.get('console_height', 120),
            simulation=simulation_config
        )

    @staticmethod
    def _dict_to_experiment_config(data: Dict[str, Any]) -> ExperimentConfig:
        """Convert dictionary to ExperimentConfig."""
        # Extract base config if present
        base_data = data.get('base_config', {})
        base_config = ConfigLoader._dict_to_headless_config(base_data)

        return ExperimentConfig(
            name=data.get('name', 'experiment'),
            description=data.get('description', ''),
            runs=data.get('runs', 1),
            run_duration=data.get('run_duration'),
            run_ticks=data.get('run_ticks'),
            variables=data.get('variables', {}),
            base_config=base_config,
            aggregate_results=data.get('aggregate_results', True),
            aggregate_format=data.get('aggregate_format', 'csv')
        )

    @staticmethod
    def get_default_headless_config() -> HeadlessConfig:
        """Get default headless configuration."""
        return HeadlessConfig()

    @staticmethod
    def get_default_interactive_config() -> InteractiveConfig:
        """Get default interactive configuration."""
        return InteractiveConfig()

    @staticmethod
    def create_sample_configs(output_dir: str = "configs"):
        """Create sample configuration files."""
        output_path = Path(output_dir)
        output_path.mkdir(exist_ok=True)

        # Sample headless config
        headless_config = ConfigLoader.get_default_headless_config()
        ConfigLoader.save_config(headless_config, output_path / "headless_default.json")

        # Sample interactive config
        interactive_config = ConfigLoader.get_default_interactive_config()
        ConfigLoader.save_config(interactive_config, output_path / "interactive_default.json")

        # Sample experiment config
        experiment_config = ExperimentConfig(
            name="tps_sweep",
            description="Test different TPS values",
            runs=5,
            run_duration=60.0,  # 1 minute per run
            variables={
                "simulation.default_tps": [10, 20, 40, 80, 160]
            }
        )
        ConfigLoader.save_config(experiment_config, output_path / "experiment_tps_sweep.json")

        print(f"Sample configs created in {output_path}")


# Import OutputConfig for the loader
from .simulation_config import OutputConfig, SimulationConfig