Atomizer/optimization_engine/run_optimization.py

"""
Generic Optimization Runner - Phase 3.2 Integration
===================================================

Flexible optimization runner supporting both manual and LLM modes:

**LLM Mode** (Natural Language):
    python run_optimization.py --llm "maximize displacement, ensure safety factor > 4" \\
        --prt model/part.prt --sim model/sim.sim

**Manual Mode** (JSON Config):
    python run_optimization.py --config config.json \\
        --prt model/part.prt --sim model/sim.sim

Features:
- Phase 2.7: LLM workflow analysis from natural language
- Phase 3.1: Auto-generated extractors
- Phase 2.9: Auto-generated hooks
- Phase 1: Plugin system with lifecycle hooks
- Graceful fallback if LLM generation fails

Author: Antoine Letarte
Version: 1.0.0 (Phase 3.2)
Last Updated: 2025-11-17
"""

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

# Add parent directory to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent))

from optimization_engine.llm_workflow_analyzer import LLMWorkflowAnalyzer
from optimization_engine.llm_optimization_runner import LLMOptimizationRunner
from optimization_engine.runner import OptimizationRunner
from optimization_engine.nx_updater import NXParameterUpdater
from optimization_engine.nx_solver import NXSolver

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


def print_banner(text: str):
    """Print a formatted banner."""
    print()
    print("=" * 80)
    print(f"  {text}")
    print("=" * 80)
    print()


def parse_arguments():
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser(
        description="Atomizer Optimization Runner - Phase 3.2 Integration",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:

  LLM Mode (Natural Language):
    python run_optimization.py \\
        --llm "maximize displacement, ensure safety factor > 4" \\
        --prt model/Bracket.prt \\
        --sim model/Bracket_sim1.sim \\
        --trials 20

  Manual Mode (JSON Config):
    python run_optimization.py \\
        --config config.json \\
        --prt model/Bracket.prt \\
        --sim model/Bracket_sim1.sim \\
        --trials 50

  With custom output directory:
    python run_optimization.py \\
        --llm "minimize stress" \\
        --prt model/part.prt \\
        --sim model/sim.sim \\
        --output results/my_study
        """
    )

    # Mode selection (mutually exclusive)
    mode_group = parser.add_mutually_exclusive_group(required=True)
    mode_group.add_argument(
        '--llm',
        type=str,
        help='Natural language optimization request (LLM mode)'
    )
    mode_group.add_argument(
        '--config',
        type=Path,
        help='Path to JSON configuration file (manual mode)'
    )

    # Required arguments
    parser.add_argument(
        '--prt',
        type=Path,
        required=True,
        help='Path to NX part file (.prt)'
    )
    parser.add_argument(
        '--sim',
        type=Path,
        required=True,
        help='Path to NX simulation file (.sim)'
    )

    # Optional arguments
    parser.add_argument(
        '--trials',
        type=int,
        default=20,
        help='Number of optimization trials (default: 20)'
    )
    parser.add_argument(
        '--output',
        type=Path,
        help='Output directory for results (default: ./optimization_results)'
    )
    parser.add_argument(
        '--study-name',
        type=str,
        help='Study name (default: auto-generated from timestamp)'
    )
    parser.add_argument(
        '--nastran-version',
        type=str,
        default='2412',
        help='Nastran version (default: 2412)'
    )
    parser.add_argument(
        '--api-key',
        type=str,
        help='Anthropic API key for LLM mode (uses Claude Code by default)'
    )

    return parser.parse_args()


def run_llm_mode(args) -> Dict[str, Any]:
    """
    Run optimization in LLM mode (natural language request).

    This uses the LLM workflow analyzer to parse the natural language request,
    then runs optimization with auto-generated extractors and hooks.

    Args:
        args: Parsed command-line arguments

    Returns:
        Optimization results dictionary
    """
    print_banner("LLM MODE - Natural Language Optimization")

    print(f"User Request: \"{args.llm}\"")
    print()

    # Step 1: Analyze natural language request using LLM
    print("Step 1: Analyzing request with LLM...")
    analyzer = LLMWorkflowAnalyzer(
        api_key=args.api_key,
        use_claude_code=(args.api_key is None)
    )

    try:
        llm_workflow = analyzer.analyze_request(args.llm)
        logger.info("LLM analysis complete!")
        logger.info(f"  Engineering features: {len(llm_workflow.get('engineering_features', []))}")
        logger.info(f"  Inline calculations: {len(llm_workflow.get('inline_calculations', []))}")
        logger.info(f"  Post-processing hooks: {len(llm_workflow.get('post_processing_hooks', []))}")
        print()

        # Validate LLM workflow structure
        required_fields = ['engineering_features', 'optimization']
        missing_fields = [f for f in required_fields if f not in llm_workflow]
        if missing_fields:
            raise ValueError(f"LLM workflow missing required fields: {missing_fields}")

        if 'design_variables' not in llm_workflow.get('optimization', {}):
            raise ValueError("LLM workflow optimization section missing 'design_variables'")

        logger.info("LLM workflow validation passed")

    except Exception as e:
        logger.error(f"LLM analysis failed: {e}")
        logger.error("Falling back to manual mode - please provide a config.json file")
        sys.exit(1)

    # Step 2: Create model updater and simulation runner
    print("Step 2: Setting up model updater and simulation runner...")

    updater = NXParameterUpdater(prt_file_path=args.prt)
    def model_updater(design_vars: dict):
        # Note: update_expressions() uses NX journal to update the file directly (use_nx_import=True by default)
        # so we don't need to call save() which would overwrite with stale binary content
        updater.update_expressions(design_vars)

    solver = NXSolver(nastran_version=args.nastran_version, use_journal=True)
    def simulation_runner(design_vars: dict) -> Path:
        result = solver.run_simulation(args.sim, expression_updates=design_vars)
        return result['op2_file']

    logger.info("  Model updater ready")
    logger.info("  Simulation runner ready")
    print()

    # Step 3: Initialize LLM optimization runner
    print("Step 3: Initializing LLM optimization runner...")

    # Determine output directory
    if args.output:
        output_dir = args.output
    else:
        output_dir = Path.cwd() / "optimization_results"

    # Determine study name
    if args.study_name:
        study_name = args.study_name
    else:
        study_name = f"llm_optimization_{datetime.now().strftime('%Y%m%d_%H%M%S')}"

    try:
        runner = LLMOptimizationRunner(
            llm_workflow=llm_workflow,
            model_updater=model_updater,
            simulation_runner=simulation_runner,
            study_name=study_name,
            output_dir=output_dir / study_name
        )

        logger.info(f"  Study name: {study_name}")
        logger.info(f"  Output directory: {runner.output_dir}")
        logger.info(f"  Extractors: {len(runner.extractors)}")
        logger.info(f"  Hooks: {runner.hook_manager.get_summary()['enabled_hooks']}")
        print()

    except Exception as e:
        logger.error(f"Failed to initialize LLM optimization runner: {e}")
        logger.error("This may be due to extractor generation or hook initialization failure")
        import traceback
        traceback.print_exc()
        sys.exit(1)

    # Step 4: Run optimization
    print_banner(f"RUNNING OPTIMIZATION - {args.trials} TRIALS")
    print(f"This will take several minutes...")
    print()

    start_time = datetime.now()
    results = runner.run_optimization(n_trials=args.trials)
    end_time = datetime.now()

    duration = (end_time - start_time).total_seconds()

    print()
    print_banner("OPTIMIZATION COMPLETE!")
    print(f"Duration: {duration:.1f} seconds ({duration/60:.1f} minutes)")
    print(f"Trials completed: {len(results['history'])}")
    print()
    print("Best Design Found:")
    for param, value in results['best_params'].items():
        print(f"  - {param}: {value:.4f}")
    print(f"  - Objective value: {results['best_value']:.6f}")
    print()
    print(f"Results saved to: {runner.output_dir}")
    print()

    return results


def run_manual_mode(args) -> Dict[str, Any]:
    """
    Run optimization in manual mode (JSON config file).

    NOTE: Manual mode integration is in progress (Task 1.2).
    For now, please use study-specific run_optimization.py scripts.

    Args:
        args: Parsed command-line arguments

    Returns:
        Optimization results dictionary
    """
    print_banner("MANUAL MODE - JSON Configuration")

    print(f"Configuration file: {args.config}")
    print()

    logger.warning("="*80)
    logger.warning("MANUAL MODE - Phase 3.2 Task 1.2 (In Progress)")
    logger.warning("="*80)
    logger.warning("")
    logger.warning("The unified runner's manual mode is currently under development.")
    logger.warning("")
    logger.warning("For manual JSON-based optimization, please use:")
    logger.warning("  - Study-specific run_optimization.py scripts")
    logger.warning("  - Example: studies/simple_beam_optimization/run_optimization.py")
    logger.warning("")
    logger.warning("Alternatively, use --llm mode for natural language optimization:")
    logger.warning("  python run_optimization.py --llm \"your request\" --prt ... --sim ...")
    logger.warning("")
    logger.warning("="*80)
    print()

    sys.exit(1)


def main():
    """Main entry point."""
    print_banner("ATOMIZER OPTIMIZATION RUNNER - Phase 3.2")

    # Parse arguments
    args = parse_arguments()

    # Validate file paths
    if not args.prt.exists():
        logger.error(f"Part file not found: {args.prt}")
        sys.exit(1)

    if not args.sim.exists():
        logger.error(f"Simulation file not found: {args.sim}")
        sys.exit(1)

    logger.info(f"Part file: {args.prt}")
    logger.info(f"Simulation file: {args.sim}")
    logger.info(f"Trials: {args.trials}")
    print()

    # Run appropriate mode
    try:
        if args.llm:
            results = run_llm_mode(args)
        else:
            results = run_manual_mode(args)

        print_banner("SUCCESS!")
        logger.info("Optimization completed successfully")

    except KeyboardInterrupt:
        print()
        logger.warning("Optimization interrupted by user")
        sys.exit(1)
    except Exception as e:
        print()
        logger.error(f"Optimization failed: {e}", exc_info=True)
        sys.exit(1)


if __name__ == '__main__':
    main()