Atomizer/tests/test_journal_optimization.py

"""
Test: Complete Optimization with Journal-Based NX Solver

This tests the complete workflow:
1. Update model parameters in .prt
2. Solve via journal (using running NX GUI)
3. Extract results from OP2
4. Run optimization loop

REQUIREMENTS:
- Simcenter3D must be open (but no files need to be loaded)
- test_env conda environment activated
"""

from pathlib import Path
import sys

# 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.runner import OptimizationRunner
from optimization_engine.nx_updater import update_nx_model
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 for the simulation runner
_current_design_vars = {}


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

    Note: We no longer directly update the .prt file here.
    Instead, design variables are passed to the journal which applies them in NX.
    """
    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.

    This connects to the running Simcenter3D GUI and:
    1. Opens the .sim file
    2. Applies expression updates in the journal
    3. Updates geometry and FEM
    4. Solves the simulation
    5. Returns path to .op2 file
    """
    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:
        # Run solver via journal (connects to running NX GUI)
        # Pass expression updates directly to the journal
        op2_file = run_nx_simulation(
            sim_file=sim_file,
            nastran_version="2412",
            timeout=300,  # 5 minute timeout
            cleanup=True,  # Clean up temp files
            use_journal=True,  # Use journal mode (requires NX GUI open)
            expression_updates=_current_design_vars  # Pass design vars to journal
        )

        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("JOURNAL-BASED OPTIMIZATION TEST")
    print("="*60)
    print("\nREQUIREMENTS:")
    print("- Simcenter3D must be OPEN (no files need to be loaded)")
    print("- Will run 50 optimization trials (~3-4 minutes)")
    print("- Strategy: 20 random trials (exploration) + 30 TPE trials (exploitation)")
    print("- Each trial: update params -> solve via journal -> extract results")
    print("="*60)

    response = input("\nIs Simcenter3D open? (yes/no): ")
    if response.lower() not in ['yes', 'y']:
        print("Please open Simcenter3D and try again.")
        sys.exit(0)

    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,  # Journal-based solver!
        result_extractors={
            'stress_extractor': stress_extractor,
            'displacement_extractor': displacement_extractor
        }
    )

    # Use the configured number of trials (50 by default)
    n_trials = runner.config['optimization_settings']['n_trials']

    # Check for existing studies
    existing_studies = runner.list_studies()

    print("\n" + "="*60)
    print("STUDY MANAGEMENT")
    print("="*60)

    if existing_studies:
        print(f"\nFound {len(existing_studies)} existing studies:")
        for study in existing_studies:
            print(f"  - {study['study_name']}: {study.get('total_trials', 0)} trials")

        print("\nOptions:")
        print("1. Create NEW study (fresh start)")
        print("2. RESUME existing study (add more trials)")
        choice = input("\nChoose option (1 or 2): ").strip()

        if choice == '2':
            # Resume existing study
            if len(existing_studies) == 1:
                study_name = existing_studies[0]['study_name']
                print(f"\nResuming study: {study_name}")
            else:
                print("\nAvailable studies:")
                for i, study in enumerate(existing_studies):
                    print(f"{i+1}. {study['study_name']}")
                study_idx = int(input("Select study number: ")) - 1
                study_name = existing_studies[study_idx]['study_name']

            resume_mode = True
        else:
            # New study
            study_name = input("\nEnter study name (default: bracket_stress_opt): ").strip()
            if not study_name:
                study_name = "bracket_stress_opt"
            resume_mode = False
    else:
        print("\nNo existing studies found. Creating new study.")
        study_name = input("\nEnter study name (default: bracket_stress_opt): ").strip()
        if not study_name:
            study_name = "bracket_stress_opt"
        resume_mode = False

    print("\n" + "="*60)
    if resume_mode:
        print(f"RESUMING STUDY: {study_name}")
        print(f"Adding {n_trials} additional trials")
    else:
        print(f"STARTING NEW STUDY: {study_name}")
        print(f"Running {n_trials} trials")
    print("="*60)
    print("Objective: Minimize max von Mises stress")
    print("Constraint: Max displacement <= 1.0 mm")
    print("Solver: Journal-based (using running NX GUI)")
    print(f"Sampler: TPE (20 random startup + {n_trials-20} TPE)")
    print("="*60)

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

        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}")

        print(f"\nResults saved to: {runner.output_dir}")
        print("\nCheck history.csv to see optimization progress!")

    except Exception as e:
        print(f"\n{'='*60}")
        print("ERROR DURING OPTIMIZATION")
        print("="*60)
        print(f"{e}")
        import traceback
        traceback.print_exc()
        print("\nMake sure:")
        print("  - Simcenter3D is open and running")
        print("  - .sim file is valid and solvable")
        print("  - No other processes are locking the files")