Atomizer/tests/test_optimization_setup_wizard.py

"""
Test Optimization Setup Wizard - Phase 3.3

This test demonstrates the complete validation workflow:
1. Introspect model for expressions
2. Run baseline simulation
3. Introspect OP2 for element types
4. Validate complete pipeline
5. Report readiness for optimization

This prevents wasted time on optimizations that will fail!
"""

import sys
from pathlib import Path

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

from optimization_engine.optimization_setup_wizard import OptimizationSetupWizard


if __name__ == '__main__':
    print("=" * 80)
    print("Phase 3.3: Optimization Setup Wizard Test")
    print("=" * 80)
    print()
    print("This wizard validates the optimization pipeline BEFORE running trials.")
    print("It will:")
    print("  1. Introspect NX model for available expressions")
    print("  2. Run baseline simulation to generate OP2")
    print("  3. Introspect OP2 to detect element types and results")
    print("  4. Validate extraction/calculation/hook pipeline")
    print("  5. Report readiness for optimization")
    print()

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

    if not prt_file.exists():
        print(f"ERROR: Part file not found: {prt_file}")
        sys.exit(1)

    if not sim_file.exists():
        print(f"ERROR: Simulation file not found: {sim_file}")
        sys.exit(1)

    print(f"Part file: {prt_file}")
    print(f"Sim file: {sim_file}")
    print()

    try:
        # Initialize wizard
        print("Initializing Optimization Setup Wizard...")
        wizard = OptimizationSetupWizard(prt_file, sim_file)
        print("  Wizard initialized!")
        print()

        # Define optimization goal
        user_goal = "Maximize displacement while keeping stress below yield/4 (safety factor >= 4.0)"

        # Custom workflow with safety factor constraint
        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'  # Will be auto-detected
                    }
                }
            ],
            '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
        }

        # Run complete validation
        print("="  * 80)
        print("RUNNING COMPLETE VALIDATION WORKFLOW")
        print("=" * 80)
        print()

        success, validation_results = wizard.run_complete_validation(
            user_goal=user_goal,
            llm_workflow=llm_workflow
        )

        print()
        print("=" * 80)
        print("VALIDATION COMPLETE")
        print("=" * 80)
        print()

        if success:
            print("[SUCCESS] ALL VALIDATIONS PASSED!")
            print()
            print("Pipeline is ready for optimization. Key findings:")
            print()

            # Show model expressions
            if wizard.model_info:
                print(f"Model Expressions ({len(wizard.model_info.expressions)}):")
                for name, info in wizard.model_info.expressions.items():
                    print(f"  - {name}: {info['value']} {info['units']}")
                print()

            # Show OP2 contents
            if wizard.op2_info:
                print(f"OP2 Contents:")
                print(f"  - Element types: {', '.join(wizard.op2_info.element_types)}")
                print(f"  - Result types: {', '.join(wizard.op2_info.result_types)}")
                print(f"  - Subcases: {wizard.op2_info.subcases}")
                print(f"  - Nodes: {wizard.op2_info.node_count}")
                print(f"  - Elements: {wizard.op2_info.element_count}")
                print()

            # Show validation summary
            print(f"Validation Summary:")
            for result in validation_results:
                if result.success:
                    print(f"  {result.message}")
            print()

            print("[OK] You can now start the optimization with confidence!")
            print("     The system has verified that all pipeline components work correctly.")

        else:
            print("[FAILED] VALIDATION FAILED!")
            print()
            print("Issues found:")
            for result in validation_results:
                if not result.success:
                    print(f"  [ERROR] {result.message}")
            print()
            print("Please fix these issues before starting optimization.")
            sys.exit(1)

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