Atomizer/tests/interactive_optimization_setup.py

"""
Interactive Optimization Setup - Phase 3.3

This script simulates a real Atomizer setup session where:
- User provides their optimization goal in natural language
- Atomizer LLM Assistant helps configure the optimization
- Wizard validates the setup before running
- User confirms and starts optimization

This is how Atomizer SHOULD work in production!
"""

import sys
from pathlib import Path

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

from optimization_engine.optimization_setup_wizard import OptimizationSetupWizard
from optimization_engine.llm_optimization_runner import LLMOptimizationRunner
from optimization_engine.nx_solver import NXSolver
from optimization_engine.nx_updater import NXParameterUpdater


def print_section(title: str):
    """Print a section header."""
    print()
    print("=" * 80)
    print(f"  {title}")
    print("=" * 80)
    print()


def print_assistant(message: str):
    """Print Atomizer LLM Assistant message."""
    print(f"\n[Atomizer] {message}\n")


def print_user(message: str):
    """Print user message."""
    print(f"\n[User] {message}\n")


def print_wizard(message: str):
    """Print wizard message."""
    print(f"[Wizard] {message}")


def main():
    print_section("ATOMIZER - INTERACTIVE OPTIMIZATION SETUP")

    print_assistant("Welcome to Atomizer! I'll help you set up your optimization.")
    print_assistant("First, I need to know about your model files.")

    # File paths
    prt_file = Path("tests/Bracket.prt")
    sim_file = Path("tests/Bracket_sim1.sim")

    if not prt_file.exists() or not sim_file.exists():
        print("[ERROR] Test files not found!")
        sys.exit(1)

    print_user(f"I have a bracket model:")
    print(f"        - Part file: {prt_file}")
    print(f"        - Simulation file: {sim_file}")

    print_assistant("Great! Let me initialize the Setup Wizard to analyze your model...")

    # Initialize wizard
    wizard = OptimizationSetupWizard(prt_file, sim_file)

    print_section("STEP 1: MODEL INTROSPECTION")

    print_assistant("I'm reading your NX model to find available design parameters...")

    # Introspect model
    model_info = wizard.introspect_model()

    print_assistant(f"Found {len(model_info.expressions)} expressions in your model:")
    for name, info in model_info.expressions.items():
        print(f"           - {name}: {info['value']} {info['units']}")

    print_assistant("Which parameters would you like to use as design variables?")

    print_user("I want to optimize tip_thickness and support_angle")

    print_assistant("Perfect! Now, what's your optimization goal?")

    print_user("I want to maximize displacement while keeping stress below")
    print("      a safety factor of 4. The material is Aluminum 6061-T6.")

    user_goal = "Maximize displacement while keeping stress below safety factor of 4 (Aluminum 6061-T6, yield=276 MPa)"

    print_section("STEP 2: BASELINE SIMULATION")

    print_assistant("To validate your setup, I need to run ONE baseline simulation.")
    print_assistant("This will generate an OP2 file that I can analyze to ensure")
    print_assistant("the extraction pipeline will work correctly.")
    print_assistant("")
    print_assistant("Running baseline simulation with current parameter values...")

    # Run baseline
    baseline_op2 = wizard.run_baseline_simulation()

    print_assistant(f"Baseline simulation complete! OP2 file: {baseline_op2.name}")

    print_section("STEP 3: OP2 INTROSPECTION")

    print_assistant("Now I'll analyze the OP2 file to see what's actually in there...")

    # Introspect OP2
    op2_info = wizard.introspect_op2()

    print_assistant("Here's what I found in your OP2 file:")
    print(f"           - Element types with stress: {', '.join(op2_info.element_types)}")
    print(f"           - Available results: {', '.join(op2_info.result_types)}")
    print(f"           - Number of elements: {op2_info.element_count}")
    print(f"           - Number of nodes: {op2_info.node_count}")

    print_section("STEP 4: LLM-GUIDED CONFIGURATION")

    print_assistant("Based on your goal and the OP2 contents, here's what I recommend:")
    print_assistant("")
    print_assistant("OBJECTIVE:")
    print_assistant("  - Maximize displacement (minimize negative displacement)")
    print_assistant("")
    print_assistant("EXTRACTIONS:")
    print_assistant("  - Extract displacement from OP2")
    print_assistant(f"  - Extract stress from {op2_info.element_types[0]} elements")
    print_assistant("    (I detected these element types in your model)")
    print_assistant("")
    print_assistant("CALCULATIONS:")
    print_assistant("  - Calculate safety factor: SF = 276 MPa / max_stress")
    print_assistant("  - Negate displacement for minimization")
    print_assistant("")
    print_assistant("CONSTRAINT:")
    print_assistant("  - Enforce SF >= 4.0 with penalty")
    print_assistant("")
    print_assistant("DESIGN VARIABLES:")
    print_assistant(f"  - tip_thickness: {model_info.expressions['tip_thickness']['value']} mm (suggest range: 15-25 mm)")
    print_assistant(f"  - support_angle: {model_info.expressions['support_angle']['value']} degrees (suggest range: 20-40 deg)")

    print_user("That looks good! Let's use those ranges.")

    # Build configuration
    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': f'Extract von Mises stress from {op2_info.element_types[0]} elements',
                'params': {
                    'result_type': 'stress',
                    'element_type': op2_info.element_types[0].lower()  # Use detected element type!
                }
            }
        ],
        'inline_calculations': [
            {
                'action': 'calculate_safety_factor',
                'params': {
                    'input': 'max_von_mises',
                    'yield_strength': 276.0,
                    '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',
            'design_variables': [
                {
                    'parameter': 'tip_thickness',
                    'min': 15.0,
                    'max': 25.0,
                    'units': 'mm'
                },
                {
                    'parameter': 'support_angle',
                    'min': 20.0,
                    'max': 40.0,
                    'units': 'degrees'
                }
            ]
        }
    }

    print_section("STEP 5: PIPELINE VALIDATION (DRY RUN)")

    print_assistant("Before running 20-30 optimization trials, let me validate that")
    print_assistant("EVERYTHING works correctly with your baseline OP2 file...")
    print_assistant("")
    print_assistant("Running dry-run validation...")

    # Validate pipeline
    validation_results = wizard.validate_pipeline(llm_workflow)

    # Check results
    all_passed = all(r.success for r in validation_results)

    print()
    print_wizard("VALIDATION RESULTS:")
    for result in validation_results:
        status = "[OK]" if result.success else "[FAIL]"
        print_wizard(f"  {status} {result.component}: {result.message.split(':')[-1].strip()}")

    print()

    if all_passed:
        print_section("VALIDATION PASSED!")

        print_assistant("Excellent! All pipeline components validated successfully.")
        print_assistant("Your optimization is ready to run!")
        print_assistant("")
        print_assistant("Summary:")
        print(f"           - Design variables: tip_thickness (15-25mm), support_angle (20-40deg)")
        print(f"           - Objective: Maximize displacement")
        print(f"           - Constraint: Safety factor >= 4.0")
        print(f"           - Material: Aluminum 6061-T6 (Yield = 276 MPa)")
        print()

        print_user("Great! Let's run 10 trials to test it.")

        print_section("STEP 6: RUNNING OPTIMIZATION")

        print_assistant("Initializing optimization runner...")

        # Create updater and runner
        def create_model_updater(prt_file: Path):
            updater = NXParameterUpdater(prt_file_path=prt_file)
            def update_model(design_vars: dict):
                updater.update_expressions(design_vars)
                updater.save()
            return update_model

        def create_simulation_runner(sim_file: Path):
            solver = NXSolver(nastran_version='2412', use_journal=True)
            def run_simulation() -> Path:
                result = solver.run_simulation(sim_file)
                return result['op2_file']
            return run_simulation

        model_updater = create_model_updater(prt_file)
        simulation_runner = create_simulation_runner(sim_file)

        # Initialize runner
        runner = LLMOptimizationRunner(
            llm_workflow=llm_workflow,
            model_updater=model_updater,
            simulation_runner=simulation_runner,
            study_name='bracket_interactive_setup'
        )

        print_assistant("Starting optimization with 10 trials...")
        print_assistant("(This will take a few minutes...)")
        print()

        # Run optimization
        results = runner.run_optimization(n_trials=10)

        print()
        print_section("OPTIMIZATION COMPLETE!")

        print_assistant("Optimization finished! Here are the results:")
        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()
        print(f"           Results saved to: {runner.output_dir}")

        # Show some trial history
        print()
        print_assistant("Trial history (last 5 trials):")
        print()
        print("           Trial | Tip(mm) | Angle(°) | Disp(mm) | Stress(MPa) | Objective")
        print("           ------|---------|----------|----------|-------------|----------")

        for trial in results['history'][-5:]:
            trial_num = trial['trial_number']
            tip = trial['design_variables']['tip_thickness']
            ang = trial['design_variables']['support_angle']
            disp = trial['results'].get('max_displacement', 0)
            stress = trial['results'].get('max_von_mises', 0)
            obj = trial['objective']
            print(f"           {trial_num:5d} | {tip:7.2f} | {ang:8.2f} | {disp:8.5f} | {stress:11.2f} | {obj:9.6f}")

        print()
        print_assistant("Success! Your optimization completed without errors.")
        print_assistant("This is exactly how Atomizer should work - validate first,")
        print_assistant("then optimize with confidence!")

    else:
        print_section("VALIDATION FAILED!")

        print_assistant("The validation found issues that need to be fixed:")
        print()

        for result in validation_results:
            if not result.success:
                print(f"           [ERROR] {result.message}")

        print()
        print_assistant("Please fix these issues before starting the optimization.")
        print_assistant("This saved you from wasting time on 20-30 failed trials!")


if __name__ == '__main__':
    main()