"""
Quick Test: Displacement-Only Optimization

Tests the pipeline with only displacement extraction (which works with your OP2).
"""

from pathlib import Path
import sys

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

from optimization_engine.runner import OptimizationRunner
from optimization_engine.nx_updater import update_nx_model
from optimization_engine.result_extractors.extractors import displacement_extractor


def bracket_model_updater(design_vars: dict):
    """Update bracket model parameters."""
    prt_file = project_root / "examples/bracket/Bracket.prt"
    print(f"\n[MODEL UPDATE] {prt_file.name}")
    for name, value in design_vars.items():
        print(f"  {name} = {value:.4f}")
    update_nx_model(prt_file, design_vars, backup=False)


def bracket_simulation_runner() -> Path:
    """Return existing OP2 (no re-solve for now)."""
    print("\n[SIMULATION] Using existing OP2")
    return project_root / "examples/bracket/bracket_sim1-solution_1.op2"


if __name__ == "__main__":
    print("="*60)
    print("DISPLACEMENT-ONLY OPTIMIZATION TEST")
    print("="*60)

    config_path = project_root / "examples/bracket/optimization_config_displacement_only.json"

    runner = OptimizationRunner(
        config_path=config_path,
        model_updater=bracket_model_updater,
        simulation_runner=bracket_simulation_runner,
        result_extractors={'displacement_extractor': displacement_extractor}
    )

    # Run 3 trials just to test
    runner.config['optimization_settings']['n_trials'] = 3

    print("\nRunning 3 test trials...")
    print("="*60)

    try:
        study = runner.run(study_name="displacement_test")

        print("\n" + "="*60)
        print("SUCCESS! Pipeline works!")
        print("="*60)
        print(f"Best displacement: {study.best_value:.6f} mm")
        print(f"Best parameters: {study.best_params}")
        print(f"\nResults in: {runner.output_dir}")

    except Exception as e:
        print(f"\nERROR: {e}")
        import traceback
        traceback.print_exc()