"""
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
5. Saves results in study directory

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
import json
from pathlib import Path
from datetime import datetime

# Add parent directories to path
sys.path.insert(0, str(Path(__file__).parent.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 save_results(results: dict, study_dir: Path):
    """Save optimization results to study directory."""
    results_dir = study_dir / "results"
    results_dir.mkdir(exist_ok=True)

    # Save complete history
    history_file = results_dir / "optimization_history.json"
    with open(history_file, 'w') as f:
        json.dump(results['history'], f, indent=2, default=str)

    # Save best design
    best_design = {
        'trial_number': results['best_trial_number'],
        'parameters': results['best_params'],
        'objective_value': results['best_value'],
        'timestamp': datetime.now().isoformat()
    }

    best_trial = results['history'][results['best_trial_number']]
    best_design['results'] = best_trial['results']
    best_design['calculations'] = best_trial['calculations']

    best_file = results_dir / "best_design.json"
    with open(best_file, 'w') as f:
        json.dump(best_design, f, indent=2, default=str)

    # Generate markdown report
    report_file = results_dir / "optimization_report.md"
    with open(report_file, 'w') as f:
        f.write("# Bracket Displacement Maximization - Optimization Report\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("- **Design Variables**:\n")
        f.write("  - tip_thickness: 15-25 mm\n")
        f.write("  - support_angle: 20-40 degrees\n\n")

        f.write("## Best Design\n\n")
        f.write(f"- **Trial**: {results['best_trial_number']}\n")
        f.write(f"- **tip_thickness**: {results['best_params']['tip_thickness']:.3f} mm\n")
        f.write(f"- **support_angle**: {results['best_params']['support_angle']:.3f} degrees\n")
        f.write(f"- **Objective value**: {results['best_value']:.6f}\n\n")

        best_results = best_trial['results']
        best_calcs = best_trial['calculations']

        f.write("## Performance\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**: {'✓ SATISFIED' if best_calcs.get('safety_factor', 0) >= 4.0 else '✗ VIOLATED'}\n\n")

        f.write("## Trial History\n\n")
        f.write("| Trial | Tip (mm) | Angle (°) | Disp (mm) | Stress (MPa) | SF | Objective |\n")
        f.write("|-------|----------|-----------|-----------|--------------|----|-----------|\n")

        for trial in results['history']:
            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)
            sf = trial['calculations'].get('safety_factor', 0)
            obj = trial['objective']
            f.write(f"| {num} | {tip:.2f} | {ang:.2f} | {disp:.6f} | {stress:.2f} | {sf:.2f} | {obj:.6f} |\n")

    return history_file, best_file, report_file


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 - USE STUDY DIRECTORY
    study_dir = Path(__file__).parent
    prt_file = study_dir / "model" / "Bracket.prt"
    sim_file = study_dir / "model" / "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(f"Study directory: {study_dir}")
    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 - UPDATE MODEL IN STUDY DIRECTORY
    updater = NXParameterUpdater(prt_file_path=prt_file)
    def model_updater(design_vars: dict):
        updater.update_expressions(design_vars)
        updater.save()

    # Simulation runner - RUN SIMULATIONS IN STUDY DIRECTORY
    solver = NXSolver(nastran_version='2412', use_journal=True)
    def simulation_runner(design_vars: dict) -> Path:
        # Pass expression values to NX journal so it can update geometry
        result = solver.run_simulation(sim_file, expression_updates=design_vars)
        return result['op2_file']

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

    print("Initializing LLM optimization runner...")

    # Save results in study/results directory
    runner = LLMOptimizationRunner(
        llm_workflow=llm_workflow,
        model_updater=model_updater,
        simulation_runner=simulation_runner,
        study_name='bracket_displacement_maximizing',
        output_dir=study_dir / "results"
    )

    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()

    # Save results
    print("Saving results...")
    history_file, best_file, report_file = save_results(results, study_dir)
    print(f"  [OK] History: {history_file.name}")
    print(f"  [OK] Best design: {best_file.name}")
    print(f"  [OK] Report: {report_file.name}")
    print()

    print_section("STUDY COMPLETE!")
    print("Phase 3.3 Optimization Setup Wizard successfully guided the")
    print("complete optimization from setup through execution!")
    print()
    print(f"Study directory: {study_dir}")
    print(f"Results directory: {study_dir / 'results'}")
    print()


if __name__ == '__main__':
    main()