Atomizer/optimization_engine/plugins/pre_solve/detailed_logger.py

"""
Detailed Logger Plugin

Logs comprehensive information about each optimization iteration to a file.
Creates a detailed trace of all steps for debugging and analysis.
"""

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

logger = logging.getLogger(__name__)


def detailed_iteration_logger(context: Dict[str, Any]) -> Optional[Dict[str, Any]]:
    """
    Log detailed information about the current trial to a timestamped log file.

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

    Returns:
        Dict with log file path
    """
    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'))

    # Create logs subdirectory within the study results
    log_dir = output_dir / 'trial_logs'
    log_dir.mkdir(parents=True, exist_ok=True)

    # Create trial-specific log file
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    log_file = log_dir / f'trial_{trial_num:03d}_{timestamp}.log'

    with open(log_file, 'w') as f:
        f.write("=" * 80 + "\n")
        f.write(f"OPTIMIZATION ITERATION LOG - Trial {trial_num}\n")
        f.write("=" * 80 + "\n")
        f.write(f"Timestamp: {datetime.now().isoformat()}\n")
        f.write(f"Output Directory: {output_dir}\n")
        f.write(f"Simulation File: {sim_file}\n")
        f.write("\n")

        f.write("-" * 80 + "\n")
        f.write("DESIGN VARIABLES\n")
        f.write("-" * 80 + "\n")
        for var_name, var_value in design_vars.items():
            f.write(f"  {var_name:30s} = {var_value:12.4f}\n")
        f.write("\n")

        f.write("-" * 80 + "\n")
        f.write("OPTIMIZATION CONFIGURATION\n")
        f.write("-" * 80 + "\n")
        config = context.get('config', {})

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

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

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

        f.write("-" * 80 + "\n")
        f.write("EXECUTION TIMELINE\n")
        f.write("-" * 80 + "\n")
        f.write(f"[{datetime.now().strftime('%H:%M:%S')}] PRE_SOLVE: Trial {trial_num} starting\n")
        f.write(f"[{datetime.now().strftime('%H:%M:%S')}] Design variables prepared\n")
        f.write(f"[{datetime.now().strftime('%H:%M:%S')}] Waiting for model update...\n")
        f.write("\n")

        f.write("-" * 80 + "\n")
        f.write("NOTES\n")
        f.write("-" * 80 + "\n")
        f.write("This log will be updated by subsequent hooks during the optimization.\n")
        f.write("Check post_solve and post_extraction logs for complete results.\n")
        f.write("\n")

    logger.info(f"Trial {trial_num} log created: {log_file}")

    return {
        'log_file': str(log_file),
        'trial_number': trial_num,
        'logged': True
    }


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=detailed_iteration_logger,
        description='Create detailed log file for each trial',
        name='detailed_logger',
        priority=5  # Run very early to capture everything
    )