Atomizer/tests/test_bracket_full_optimization.py

"""
Full End-to-End Bracket Optimization Test - Phase 3.2

This test demonstrates the complete LLM-enhanced optimization workflow:
1. LLM workflow configuration
2. Automatic extractor generation (displacement + stress)
3. Inline calculation generation (safety factor)
4. Manual safety factor constraint hook
5. Real NX simulation execution
6. OP2 result extraction
7. Constraint checking
8. Optimization with Optuna
9. Report generation

Objective: Maximize displacement
Constraint: Safety factor >= 4.0 (stress < yield/4)
Material: Aluminum 6061-T6 (Yield = 276 MPa)
Design Variable: wall_thickness (3-8mm)
"""

import sys
from pathlib import Path
from datetime import datetime
import json

sys.path.insert(0, str(Path(__file__).parent.parent))

from optimization_engine.future.llm_optimization_runner import LLMOptimizationRunner
from optimization_engine.nx.solver import NXSolver
from optimization_engine.nx.updater import NXParameterUpdater

# LLM workflow for bracket optimization
llm_workflow = {
    'engineering_features': [
        {
            'action': 'extract_displacement',
            'domain': 'result_extraction',
            'description': 'Extract displacement results from OP2 file',
            'params': {'result_type': 'displacement'}
        },
        {
            'action': 'extract_solid_stress',
            'domain': 'result_extraction',
            'description': 'Extract von Mises stress from solid elements',
            'params': {
                'result_type': 'stress',
                'element_type': 'chexa'  # Bracket uses CHEXA elements, not CTETRA
            }
        }
    ],
    'inline_calculations': [
        {
            'action': 'calculate_safety_factor',
            'params': {
                'input': 'max_von_mises',
                'yield_strength': 276.0,  # MPa for Aluminum 6061-T6
                'operation': 'divide'
            },
            'code_hint': 'safety_factor = 276.0 / max_von_mises'
        },
        {
            'action': 'negate_displacement',
            'params': {
                'input': 'max_displacement',
                'operation': 'negate'
            },
            'code_hint': 'neg_displacement = -max_displacement'
        }
    ],
    'post_processing_hooks': [],  # Using manual safety_factor_constraint hook
    'optimization': {
        'algorithm': 'TPE',
        'direction': 'minimize',  # Minimize neg_displacement = maximize displacement
        'design_variables': [
            {
                'parameter': 'tip_thickness',
                'min': 15.0,
                'max': 25.0,
                'units': 'mm'
            },
            {
                'parameter': 'support_angle',
                'min': 20.0,
                'max': 40.0,
                'units': 'degrees'
            }
        ]
    }
}


def create_model_updater(prt_file: Path):
    """Create model updater for NX."""
    updater = NXParameterUpdater(prt_file_path=prt_file)

    def update_model(design_vars: dict):
        """Update NX model with design variables."""
        updater.update_expressions(design_vars)
        updater.save()  # Save changes to disk

    return update_model


def create_simulation_runner(sim_file: Path):
    """Create simulation runner for NX."""
    solver = NXSolver(nastran_version='2412', use_journal=True)

    def run_simulation() -> Path:
        """Run NX simulation and return OP2 file path."""
        result = solver.run_simulation(sim_file)
        return result['op2_file']  # Extract Path from result dict

    return run_simulation


def generate_report(results: dict, output_dir: Path):
    """Generate optimization report with history."""
    report_file = output_dir / "optimization_report.md"

    best_params = results['best_params']
    best_value = results['best_value']
    history = results['history']

    with open(report_file, 'w') as f:
        f.write("# Bracket Optimization Report - Phase 3.2\n\n")
        f.write(f"**Generated**: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n")

        f.write("## Problem Definition\n\n")
        f.write("- **Objective**: Maximize displacement\n")
        f.write("- **Constraint**: Safety factor >= 4.0\n")
        f.write("- **Material**: Aluminum 6061-T6 (Yield = 276 MPa)\n")
        f.write("- **Allowable Stress**: 69 MPa (276/4)\n")
        f.write("- **Design Variables**:\n")
        f.write("  - tip_thickness (15-25 mm)\n")
        f.write("  - support_angle (20-40 degrees)\n\n")

        f.write("## Optimization Results\n\n")
        f.write(f"- **Best tip_thickness**: {best_params['tip_thickness']:.3f} mm\n")
        f.write(f"- **Best support_angle**: {best_params['support_angle']:.3f} degrees\n")
        f.write(f"- **Best objective value**: {best_value:.6f}\n")
        f.write(f"- **Total trials**: {len(history)}\n\n")

        # Find best trial details
        best_trial = history[results['best_trial_number']]
        best_results = best_trial['results']
        best_calcs = best_trial['calculations']

        f.write("## Best Design Details\n\n")
        f.write(f"- **Max displacement**: {best_results.get('max_displacement', 0):.6f} mm\n")
        f.write(f"- **Max stress**: {best_results.get('max_von_mises', 0):.3f} MPa\n")
        f.write(f"- **Safety factor**: {best_calcs.get('safety_factor', 0):.3f}\n")
        f.write(f"- **Constraint status**: {'SATISFIED' if best_calcs.get('safety_factor', 0) >= 4.0 else 'VIOLATED'}\n\n")

        f.write("## Optimization History\n\n")
        f.write("| Trial | Tip Thick (mm) | Support Angle (°) | Displacement (mm) | Stress (MPa) | Safety Factor | Constraint | Objective |\n")
        f.write("|-------|----------------|-------------------|-------------------|--------------|---------------|------------|-----------|\n")

        for trial in history:
            trial_num = trial['trial_number']
            tip_thick = trial['design_variables']['tip_thickness']
            support_ang = trial['design_variables']['support_angle']
            disp = trial['results'].get('max_displacement', 0)
            stress = trial['results'].get('max_von_mises', 0)
            sf = trial['calculations'].get('safety_factor', 0)
            constraint = 'OK' if sf >= 4.0 else 'FAIL'
            obj = trial['objective']

            f.write(f"| {trial_num} | {tip_thick:.3f} | {support_ang:.3f} | {disp:.6f} | {stress:.3f} | {sf:.3f} | {constraint} | {obj:.6f} |\n")

        f.write("\n## LLM-Enhanced Workflow\n\n")
        f.write("This optimization was run using the Phase 3.2 LLM-enhanced runner with:\n\n")
        f.write("- **Automatic extractor generation**: Displacement + Stress (Phase 3.1)\n")
        f.write("- **Inline calculations**: Safety factor + Objective negation (Phase 2.8)\n")
        f.write("- **Manual constraint hook**: Safety factor constraint (demonstrates flexibility)\n")
        f.write("- **Real NX simulation**: Journal-based solver execution\n")
        f.write("- **Optuna optimization**: TPE sampler\n\n")

        f.write("---\n\n")
        f.write("*Generated by Atomizer Phase 3.2 - LLM-Enhanced Optimization Framework*\n")

    print(f"\nReport saved to: {report_file}")
    return report_file


if __name__ == '__main__':
    print("=" * 80)
    print("Phase 3.2: Full End-to-End Bracket Optimization Test")
    print("=" * 80)
    print()
    print("Problem:")
    print("  Maximize displacement while maintaining safety factor >= 4.0")
    print("  Material: Aluminum 6061-T6 (Yield = 276 MPa)")
    print("  Design Variables:")
    print("    - tip_thickness (15-25 mm)")
    print("    - support_angle (20-40 degrees)")
    print()

    # Configuration
    sim_file = Path("tests/Bracket_sim1.sim")
    prt_file = Path("tests/Bracket.prt")
    n_trials = 5  # Start with 5 trials for testing

    if not sim_file.exists():
        print(f"ERROR: Simulation file not found: {sim_file}")
        print("Please ensure the bracket .sim file is available")
        sys.exit(1)

    if not prt_file.exists():
        print(f"ERROR: Part file not found: {prt_file}")
        print("Please ensure the bracket .prt file is available")
        sys.exit(1)

    print(f"Simulation file: {sim_file}")
    print(f"Part file: {prt_file}")
    print(f"Number of trials: {n_trials}")
    print()

    try:
        # Create model updater and simulation runner
        print("Setting up NX integration...")
        model_updater = create_model_updater(prt_file)
        simulation_runner = create_simulation_runner(sim_file)
        print("  NX updater: OK")
        print("  NX solver: OK")
        print()

        # Initialize LLM optimization runner
        print("Initializing LLM-enhanced optimization runner...")
        runner = LLMOptimizationRunner(
            llm_workflow=llm_workflow,
            model_updater=model_updater,
            simulation_runner=simulation_runner,
            study_name='bracket_maximize_disp_sf4'
        )

        print(f"  Extractors generated: {len(runner.extractors)}")
        for ext in runner.extractors:
            print(f"    - {ext}")
        print(f"  Inline calculations: {len(runner.inline_code)}")
        hook_summary = runner.hook_manager.get_summary()
        print(f"  Hooks loaded: {hook_summary['enabled_hooks']}")
        print("    (Including manual safety_factor_constraint hook)")
        print()

        # Run optimization
        print("=" * 80)
        print("Starting Optimization...")
        print("=" * 80)
        print()

        results = runner.run_optimization(n_trials=n_trials)

        print()
        print("=" * 80)
        print("Optimization Complete!")
        print("=" * 80)
        print()
        print(f"Best tip_thickness: {results['best_params']['tip_thickness']:.3f} mm")
        print(f"Best support_angle: {results['best_params']['support_angle']:.3f} degrees")
        print(f"Best objective value: {results['best_value']:.6f}")
        print()

        # Generate report
        print("Generating optimization report...")
        report_file = generate_report(results, runner.output_dir)

        print()
        print("=" * 80)
        print("Test Complete!")
        print("=" * 80)
        print(f"Results directory: {runner.output_dir}")
        print(f"Report: {report_file}")

    except Exception as e:
        print(f"\nERROR during optimization: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)