Atomizer/optimization_engine/plugins/pre_solve/optimization_logger.py

"""
Optimization-Level Logger Hook

Creates a high-level optimization log file that tracks the overall progress
across all trials. This complements the detailed per-trial logs.

Hook Point: pre_solve
"""

from pathlib import Path
from datetime import datetime
from typing import Dict, Any, Optional
import logging

logger = logging.getLogger(__name__)


def log_optimization_progress(context: Dict[str, Any]) -> Optional[Dict[str, Any]]:
    """
    Log high-level optimization progress to optimization.log.

    This hook creates/appends to a main optimization log file that shows:
    - Trial start with design variables
    - High-level progress tracking
    - Easy-to-scan overview of the optimization run

    Args:
        context: Hook context containing:
            - trial_number: Current trial number
            - design_variables: Dict of variable values
            - sim_file: Path to simulation file
            - config: Full optimization configuration

    Returns:
        None (logging only)
    """
    trial_num = context.get('trial_number', '?')
    design_vars = context.get('design_variables', {})
    sim_file = context.get('sim_file', 'unknown')
    config = context.get('config', {})

    # Get the output directory from context (passed by runner)
    output_dir = Path(context.get('output_dir', 'optimization_results'))

    # Main optimization log file
    log_file = output_dir / 'optimization.log'

    # Create header on first trial
    if trial_num == 0:
        output_dir.mkdir(parents=True, exist_ok=True)
        with open(log_file, 'w') as f:
            f.write("=" * 100 + "\n")
            f.write(f"OPTIMIZATION RUN - Started {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
            f.write("=" * 100 + "\n")
            f.write(f"Simulation File: {sim_file}\n")
            f.write(f"Output Directory: {output_dir}\n")

            # Optimization settings
            opt_settings = config.get('optimization_settings', {})
            f.write(f"\nOptimization Settings:\n")
            f.write(f"  Total Trials: {opt_settings.get('n_trials', 'unknown')}\n")
            f.write(f"  Sampler: {opt_settings.get('sampler', 'unknown')}\n")
            f.write(f"  Startup Trials: {opt_settings.get('n_startup_trials', 'unknown')}\n")

            # Design variables
            design_vars_config = config.get('design_variables', [])
            f.write(f"\nDesign Variables:\n")
            for dv in design_vars_config:
                name = dv.get('name', 'unknown')
                bounds = dv.get('bounds', [])
                units = dv.get('units', '')
                f.write(f"  {name}: {bounds[0]:.2f} - {bounds[1]:.2f} {units}\n")

            # Objectives
            objectives = config.get('objectives', [])
            f.write(f"\nObjectives:\n")
            for obj in objectives:
                name = obj.get('name', 'unknown')
                direction = obj.get('direction', 'unknown')
                units = obj.get('units', '')
                f.write(f"  {name} ({direction}) [{units}]\n")

            # Constraints
            constraints = config.get('constraints', [])
            if constraints:
                f.write(f"\nConstraints:\n")
                for cons in constraints:
                    name = cons.get('name', 'unknown')
                    cons_type = cons.get('type', 'unknown')
                    limit = cons.get('limit', 'unknown')
                    units = cons.get('units', '')
                    f.write(f"  {name}: {cons_type} {limit} {units}\n")

            f.write("\n" + "=" * 100 + "\n")
            f.write("TRIAL PROGRESS\n")
            f.write("=" * 100 + "\n\n")

    # Append trial start
    with open(log_file, 'a') as f:
        timestamp = datetime.now().strftime('%H:%M:%S')
        f.write(f"[{timestamp}] Trial {trial_num:3d} START | ")

        # Write design variables in compact format
        dv_str = ", ".join([f"{name}={value:.3f}" for name, value in design_vars.items()])
        f.write(f"{dv_str}\n")

    return None


def register_hooks(hook_manager):
    """
    Register this plugin's hooks with the manager.

    This function is called automatically when the plugin is loaded.
    """
    hook_manager.register_hook(
        hook_point='pre_solve',
        function=log_optimization_progress,
        description='Create high-level optimization.log file',
        name='optimization_logger',
        priority=100  # Run early to set up log file
    )


# Hook metadata
HOOK_NAME = "optimization_logger"
HOOK_POINT = "pre_solve"
ENABLED = True
PRIORITY = 100  # Run early to set up log file