Atomizer/optimization_engine/plugins/post_solve/log_solve_complete.py

"""
Post-Solve Logger Plugin

Appends solver completion information to the trial log.
"""

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

logger = logging.getLogger(__name__)


def log_solve_complete(context: Dict[str, Any]) -> Optional[Dict[str, Any]]:
    """
    Log solver completion information to the trial log file.

    Args:
        context: Hook context containing:
            - trial_number: Current trial number
            - design_variables: Dict of variable values
            - result_path: Path to OP2 result file
            - working_dir: Current working directory
    """
    trial_num = context.get('trial_number', '?')
    result_path = context.get('result_path', 'unknown')

    # Get the output directory from context (passed by runner)
    output_dir = Path(context.get('output_dir', 'optimization_results'))
    log_dir = output_dir / 'trial_logs'
    if not log_dir.exists():
        logger.warning(f"Log directory not found: {log_dir}")
        return None

    # Find trial log file
    log_files = list(log_dir.glob(f'trial_{trial_num:03d}_*.log'))
    if not log_files:
        logger.warning(f"No log file found for trial {trial_num}")
        return None

    # Use most recent log file
    log_file = sorted(log_files)[-1]

    with open(log_file, 'a') as f:
        f.write(f"[{datetime.now().strftime('%H:%M:%S')}] POST_SOLVE: Simulation complete\n")
        f.write(f"[{datetime.now().strftime('%H:%M:%S')}] Result file: {Path(result_path).name}\n")
        f.write(f"[{datetime.now().strftime('%H:%M:%S')}] Result path: {result_path}\n")
        f.write(f"[{datetime.now().strftime('%H:%M:%S')}] Waiting for result extraction...\n")
        f.write("\n")

    return {'logged': True}


def register_hooks(hook_manager):
    """Register this plugin's hooks with the manager."""
    hook_manager.register_hook(
        hook_point='post_solve',
        function=log_solve_complete,
        description='Log solver completion to trial log',
        name='log_solve_complete',
        priority=10
    )