"""
Quick 5-Trial Test with Real NX Solves

Tests the complete optimization pipeline with hooks and logging.
"""

from pathlib import Path
import sys
from datetime import datetime

# Add project root to path so we can import atomizer_paths
sys.path.insert(0, str(Path(__file__).parent.parent))

# Use intelligent path resolution
import atomizer_paths
atomizer_paths.ensure_imports()

from optimization_engine.core.runner import OptimizationRunner
from optimization_engine.nx.solver import run_nx_simulation
from optimization_engine.result_extractors.extractors import (
    stress_extractor,
    displacement_extractor
)

# Global variable to store current design variables
_current_design_vars = {}


def bracket_model_updater(design_vars: dict):
    """Store design variables for the simulation runner."""
    global _current_design_vars
    _current_design_vars = design_vars.copy()

    print(f"\n[MODEL UPDATE] Design variables prepared")
    for name, value in design_vars.items():
        print(f"  {name} = {value:.4f}")


def bracket_simulation_runner() -> Path:
    """Run NX solver via journal on running NX GUI session."""
    global _current_design_vars
    sim_file = atomizer_paths.studies() / "bracket_stress_minimization/model/Bracket_sim1.sim"

    print("\n[SIMULATION] Running via journal on NX GUI...")
    print(f"  SIM file: {sim_file.name}")
    if _current_design_vars:
        print(f"  Expression updates: {_current_design_vars}")

    try:
        op2_file = run_nx_simulation(
            sim_file=sim_file,
            nastran_version="2412",
            timeout=300,
            cleanup=True,
            use_journal=True,
            expression_updates=_current_design_vars
        )

        print(f"[SIMULATION] Complete! Results: {op2_file.name}")
        return op2_file

    except Exception as e:
        print(f"[SIMULATION] FAILED: {e}")
        raise


if __name__ == "__main__":
    print("="*60)
    print("5-TRIAL REAL OPTIMIZATION TEST WITH LOGGING")
    print("="*60)
    print("\nThis will:")
    print("- Run 5 real NX solves (not dummy data)")
    print("- Create detailed log files in optimization_logs/")
    print("- Test all hooks (pre_solve, post_solve, post_extraction)")
    print("- Verify design variables actually change results")
    print("\nTime: ~3-5 minutes")
    print("="*60)

    input("\nPress ENTER to continue (Ctrl+C to cancel)...")

    config_path = atomizer_paths.studies() / "bracket_stress_minimization/optimization_config_stress_displacement.json"

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

    # Generate unique study name with timestamp
    study_name = f"test_5trials_{datetime.now().strftime('%Y%m%d_%H%M%S')}"

    print("\n" + "="*60)
    print("RUNNING 5-TRIAL OPTIMIZATION")
    print("="*60)
    print(f"Study name: {study_name}")
    print("="*60)

    try:
        study = runner.run(
            study_name=study_name,
            n_trials=5,
            resume=False
        )

        print("\n" + "="*60)
        print("OPTIMIZATION COMPLETE!")
        print("="*60)
        print(f"\nBest stress: {study.best_value:.2f} MPa")
        print(f"\nBest parameters:")
        for param, value in study.best_params.items():
            print(f"  {param}: {value:.4f}")

        # Show log files created
        log_dir = runner.output_dir / "trial_logs"
        if log_dir.exists():
            log_files = sorted(log_dir.glob("trial_*.log"))
            print(f"\n{len(log_files)} detailed trial logs created:")
            print(f"  {log_dir}")
            print("\nExample log file (open to see detailed iteration trace):")
            if log_files:
                print(f"  {log_files[0].name}")

        print(f"\nResults also saved to: {runner.output_dir}")

    except Exception as e:
        print(f"\n{'='*60}")
        print("ERROR DURING OPTIMIZATION")
        print("="*60)
        print(f"{e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)