Atomizer/optimization_engine/config/manager.py

"""Configuration validation and management for Atomizer studies.

This module provides schema-based validation for optimization configuration files,
ensuring consistency across all studies.

Usage:
    # In run_optimization.py
    from optimization_engine.config.manager import ConfigManager

    config_manager = ConfigManager(Path(__file__).parent / "1_setup" / "optimization_config.json")
    config_manager.load_config()

    if not config_manager.validate():
        print(config_manager.get_validation_report())
        sys.exit(1)

    # Access validated configuration
    design_vars = config_manager.get_design_variables()
    objectives = config_manager.get_objectives()
"""

import json
from pathlib import Path
from typing import Dict, List, Any, Optional

try:
    import jsonschema
    JSONSCHEMA_AVAILABLE = True
except ImportError:
    JSONSCHEMA_AVAILABLE = False
    print("Warning: jsonschema not installed. Install with: pip install jsonschema>=4.17.0")


class ConfigValidationError(Exception):
    """Raised when configuration validation fails."""
    pass


class ConfigManager:
    """Manages and validates optimization configuration files."""

    def __init__(self, config_path: Path):
        """
        Initialize ConfigManager with path to optimization_config.json.

        Args:
            config_path: Path to optimization_config.json file
        """
        self.config_path = Path(config_path)
        self.schema_path = Path(__file__).parent / "schemas" / "optimization_config_schema.json"
        self.config: Optional[Dict[str, Any]] = None
        self.validation_errors: List[str] = []

    def load_schema(self) -> Dict[str, Any]:
        """Load JSON schema for validation."""
        if not self.schema_path.exists():
            raise FileNotFoundError(f"Schema file not found: {self.schema_path}")

        with open(self.schema_path, 'r') as f:
            return json.load(f)

    def load_config(self) -> Dict[str, Any]:
        """Load configuration file."""
        if not self.config_path.exists():
            raise FileNotFoundError(f"Config file not found: {self.config_path}")

        with open(self.config_path, 'r') as f:
            self.config = json.load(f)
        return self.config

    def validate(self, strict: bool = True) -> bool:
        """
        Validate configuration against schema.

        Args:
            strict: If True, enforce all validations. If False, only warn on non-critical issues.

        Returns:
            True if valid, False otherwise
        """
        if self.config is None:
            self.load_config()

        self.validation_errors = []

        # JSON Schema validation
        if JSONSCHEMA_AVAILABLE:
            schema = self.load_schema()
            try:
                jsonschema.validate(instance=self.config, schema=schema)
            except jsonschema.ValidationError as e:
                self.validation_errors.append(f"Schema validation failed: {e.message}")
                if strict:
                    return False
        else:
            self.validation_errors.append("jsonschema not installed - schema validation skipped")

        # Custom validations
        self._validate_design_variable_bounds()
        self._validate_multi_objective_consistency()
        self._validate_file_locations()
        self._validate_extraction_consistency()

        return len(self.validation_errors) == 0

    def _validate_design_variable_bounds(self):
        """Ensure bounds are valid (min < max)."""
        for dv in self.config.get("design_variables", []):
            bounds = dv.get("bounds", [])
            if len(bounds) == 2 and bounds[0] >= bounds[1]:
                self.validation_errors.append(
                    f"Design variable '{dv.get('parameter', 'unknown')}': "
                    f"min ({bounds[0]}) must be < max ({bounds[1]})"
                )

    def _validate_multi_objective_consistency(self):
        """Validate multi-objective settings consistency."""
        n_objectives = len(self.config.get("objectives", []))
        protocol = self.config.get("optimization_settings", {}).get("protocol")
        sampler = self.config.get("optimization_settings", {}).get("sampler")

        if n_objectives > 1:
            # Multi-objective should use protocol_11 and NSGA-II
            if protocol and protocol != "protocol_11_multi_objective":
                self.validation_errors.append(
                    f"Multi-objective optimization ({n_objectives} objectives) "
                    f"should use protocol_11_multi_objective (got {protocol})"
                )
            if sampler and sampler != "NSGAIISampler":
                self.validation_errors.append(
                    f"Multi-objective optimization should use NSGAIISampler (got {sampler})"
                )
        elif n_objectives == 1:
            # Single-objective should not use NSGA-II
            if sampler == "NSGAIISampler":
                self.validation_errors.append(
                    "Single-objective optimization should not use NSGAIISampler "
                    "(use TPESampler or CmaEsSampler)"
                )

    def _validate_file_locations(self):
        """Check if config is in correct location (1_setup/)."""
        if "1_setup" not in str(self.config_path.parent):
            self.validation_errors.append(
                f"Warning: Config should be in '1_setup/' directory, "
                f"found in {self.config_path.parent}"
            )

    def _validate_extraction_consistency(self):
        """Validate extraction specifications."""
        # Check objectives have extraction specs
        for obj in self.config.get("objectives", []):
            if "extraction" not in obj:
                self.validation_errors.append(
                    f"Objective '{obj.get('name', 'unknown')}' missing extraction specification"
                )

        # Check constraints have extraction specs
        for constraint in self.config.get("constraints", []):
            if "extraction" not in constraint:
                self.validation_errors.append(
                    f"Constraint '{constraint.get('name', 'unknown')}' missing extraction specification"
                )

    def get_validation_report(self) -> str:
        """Get human-readable validation report."""
        if not self.validation_errors:
            return "[OK] Configuration is valid"

        report = "[FAIL] Configuration validation failed:\n"
        for i, error in enumerate(self.validation_errors, 1):
            report += f"  {i}. {error}\n"
        return report

    # Type-safe accessor methods

    def get_design_variables(self) -> List[Dict[str, Any]]:
        """Get design variables with validated structure."""
        if self.config is None:
            self.load_config()
        return self.config.get("design_variables", [])

    def get_objectives(self) -> List[Dict[str, Any]]:
        """Get objectives with validated structure."""
        if self.config is None:
            self.load_config()
        return self.config.get("objectives", [])

    def get_constraints(self) -> List[Dict[str, Any]]:
        """Get constraints with validated structure."""
        if self.config is None:
            self.load_config()
        return self.config.get("constraints", [])

    def get_simulation_settings(self) -> Dict[str, Any]:
        """Get simulation settings."""
        if self.config is None:
            self.load_config()
        return self.config.get("simulation", {})

    def get_optimization_settings(self) -> Dict[str, Any]:
        """Get optimization settings."""
        if self.config is None:
            self.load_config()
        return self.config.get("optimization_settings", {})


# CLI tool for validation
if __name__ == "__main__":
    import sys

    if len(sys.argv) < 2:
        print("Usage: python config_manager.py <path_to_optimization_config.json>")
        print("\nExample:")
        print("  python config_manager.py studies/drone_gimbal_arm_optimization/1_setup/optimization_config.json")
        sys.exit(1)

    config_path = Path(sys.argv[1])

    print(f"Validating configuration: {config_path}")
    print("=" * 60)

    manager = ConfigManager(config_path)

    try:
        manager.load_config()
        print("[OK] Config loaded successfully")

        is_valid = manager.validate()
        print(manager.get_validation_report())

        if is_valid:
            print("\n" + "=" * 60)
            print("Configuration Summary:")
            print(f"  Study: {manager.config.get('study_name')}")
            print(f"  Protocol: {manager.get_optimization_settings().get('protocol')}")
            print(f"  Design Variables: {len(manager.get_design_variables())}")
            print(f"  Objectives: {len(manager.get_objectives())}")
            print(f"  Constraints: {len(manager.get_constraints())}")

        sys.exit(0 if is_valid else 1)
    except Exception as e:
        print(f"[ERROR] {e}")
        sys.exit(1)