"""
Bracket Displacement Maximization Study
========================================

Complete optimization workflow using Phase 3.3 Wizard:
1. Setup wizard validates the complete pipeline
2. Auto-detects element types from OP2
3. Runs 20-trial optimization
4. Generates comprehensive report

Objective: Maximize displacement
Constraint: Safety factor >= 4.0
Material: Aluminum 6061-T6 (Yield = 276 MPa)
Design Variables: tip_thickness (15-25mm), support_angle (20-40deg)
"""

import sys
from pathlib import Path

# Add parent directories to path
sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from optimization_engine.config.setup_wizard import OptimizationSetupWizard
from optimization_engine.future.llm_optimization_runner import LLMOptimizationRunner
from optimization_engine.nx.solver import NXSolver
from optimization_engine.nx.updater import NXParameterUpdater
from datetime import datetime


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


def main():
    print_section("BRACKET DISPLACEMENT MAXIMIZATION STUDY")

    print("Study Configuration:")
    print("  - Objective: Maximize displacement")
    print("  - Constraint: 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("  - Optimization trials: 20")
    print()

    # File paths
    base_dir = Path(__file__).parent.parent.parent
    prt_file = base_dir / "tests" / "Bracket.prt"
    sim_file = base_dir / "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"Simulation file: {sim_file}")
    print()

    # =========================================================================
    # PHASE 3.3: OPTIMIZATION SETUP WIZARD
    # =========================================================================

    print_section("STEP 1: INITIALIZATION")

    print("Initializing Optimization Setup Wizard...")
    wizard = OptimizationSetupWizard(prt_file, sim_file)
    print("  [OK] Wizard initialized")
    print()

    print_section("STEP 2: MODEL INTROSPECTION")

    print("Reading NX model expressions...")
    model_info = wizard.introspect_model()

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

    print_section("STEP 3: BASELINE SIMULATION")

    print("Running baseline simulation to generate reference OP2...")
    print("(This validates that NX simulation works before optimization)")
    baseline_op2 = wizard.run_baseline_simulation()
    print(f"  [OK] Baseline OP2: {baseline_op2.name}")
    print()

    print_section("STEP 4: OP2 INTROSPECTION")

    print("Analyzing OP2 file to auto-detect element types...")
    op2_info = wizard.introspect_op2()

    print("OP2 Contents:")
    print(f"  - Element types with stress: {', '.join(op2_info.element_types)}")
    print(f"  - Available result types: {', '.join(op2_info.result_types)}")
    print(f"  - Subcases: {op2_info.subcases}")
    print(f"  - Nodes: {op2_info.node_count}")
    print(f"  - Elements: {op2_info.element_count}")
    print()

    print_section("STEP 5: WORKFLOW CONFIGURATION")

    print("Building LLM workflow with auto-detected element types...")

    # Use the FIRST detected element type (could be CHEXA, CPENTA, CTETRA, etc.)
    detected_element_type = op2_info.element_types[0].lower() if op2_info.element_types else 'ctetra'

    print(f"  Using detected element type: {detected_element_type.upper()}")
    print()

    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 {detected_element_type.upper()} elements',
                'params': {
                    'result_type': 'stress',
                    'element_type': detected_element_type  # 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
        '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'
                }
            ]
        }
    }

    print_section("STEP 6: PIPELINE VALIDATION")

    print("Validating complete pipeline with baseline OP2...")
    print("(Dry-run test of extractors, calculations, hooks, objective)")
    print()

    validation_results = wizard.validate_pipeline(llm_workflow)

    all_passed = all(r.success for r in validation_results)

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

    if not all_passed:
        print("[FAILED] Pipeline validation failed!")
        print("Fix the issues above before running optimization.")
        sys.exit(1)

    print("[SUCCESS] All pipeline components validated!")
    print()

    print_section("STEP 7: OPTIMIZATION SETUP")

    print("Creating model updater and simulation runner...")

    # Model updater
    updater = NXParameterUpdater(prt_file_path=prt_file)
    def model_updater(design_vars: dict):
        updater.update_expressions(design_vars)
        updater.save()

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

    print("  [OK] Model updater ready")
    print("  [OK] Simulation runner ready")
    print()

    print("Initializing LLM optimization runner...")
    runner = LLMOptimizationRunner(
        llm_workflow=llm_workflow,
        model_updater=model_updater,
        simulation_runner=simulation_runner,
        study_name='bracket_displacement_maximizing'
    )

    print(f"  [OK] Output directory: {runner.output_dir}")
    print(f"  [OK] Extractors generated: {len(runner.extractors)}")
    print(f"  [OK] Inline calculations: {len(runner.inline_code)}")
    hook_summary = runner.hook_manager.get_summary()
    print(f"  [OK] Hooks loaded: {hook_summary['enabled_hooks']}")
    print()

    print_section("STEP 8: RUNNING OPTIMIZATION")

    print("Starting 20-trial optimization...")
    print("(This will take several minutes)")
    print()

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

    duration = (end_time - start_time).total_seconds()

    print()
    print_section("OPTIMIZATION COMPLETE!")

    print(f"Duration: {duration:.1f} seconds ({duration/60:.1f} minutes)")
    print()
    print("Best Design Found:")
    print(f"  - tip_thickness: {results['best_params']['tip_thickness']:.3f} mm")
    print(f"  - support_angle: {results['best_params']['support_angle']:.3f} degrees")
    print(f"  - Objective value: {results['best_value']:.6f}")
    print()

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

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

    print(f"Results saved to: {runner.output_dir}")
    print()

    print_section("STUDY COMPLETE!")
    print("Phase 3.3 Optimization Setup Wizard successfully guided the")
    print("complete optimization from setup through execution!")
    print()


if __name__ == '__main__':
    main()