Atomizer/optimization_engine/nx_solver.py

"""
NX Nastran Solver Integration

Executes NX Nastran solver in batch mode for optimization loops.
Includes session management to prevent conflicts with concurrent optimizations.
"""

from pathlib import Path
from typing import Optional, Dict, Any
import subprocess
import time
import shutil
import os
from optimization_engine.nx_session_manager import NXSessionManager


class NXSolver:
    """
    Wrapper for NX Nastran batch solver execution.

    Supports:
    - Running .sim files through NX Nastran
    - Monitoring solver progress
    - Detecting completion and errors
    - Cleaning up temporary files
    - Per-iteration model copies (HEEDS-style isolation)
    """

    # Model file extensions to copy for each iteration
    MODEL_EXTENSIONS = {'.prt', '.fem', '.afm', '.sim', '.exp'}

    def __init__(
        self,
        nx_install_dir: Optional[Path] = None,
        nastran_version: str = "2412",
        timeout: int = 600,
        use_journal: bool = True,
        enable_session_management: bool = True,
        study_name: str = "default_study",
        use_iteration_folders: bool = False,
        master_model_dir: Optional[Path] = None
    ):
        """
        Initialize NX Solver.

        Args:
            nx_install_dir: Path to NX installation (auto-detected if None)
            nastran_version: NX version (e.g., "2412", "2506")
            timeout: Maximum solver time in seconds (default: 10 minutes)
            use_journal: Use NX journal for solving (recommended for licensing)
            enable_session_management: Enable session conflict prevention (default: True)
            study_name: Name of the study (used for session tracking)
            use_iteration_folders: Create Iter1, Iter2, etc. folders with fresh model copies
            master_model_dir: Source directory for master model files (required if use_iteration_folders=True)
        """
        self.nastran_version = nastran_version
        self.timeout = timeout
        self.use_journal = use_journal
        self.study_name = study_name
        self.use_iteration_folders = use_iteration_folders
        self.master_model_dir = Path(master_model_dir) if master_model_dir else None
        self._iteration_counter = 0

        # Initialize session manager
        self.session_manager = None
        if enable_session_management:
            self.session_manager = NXSessionManager(verbose=True)
            # Clean up any stale locks from crashed processes
            self.session_manager.cleanup_stale_locks()

        # Auto-detect NX installation
        if nx_install_dir is None:
            nx_install_dir = self._find_nx_installation()

        self.nx_install_dir = Path(nx_install_dir)

        # Set up solver executable
        if use_journal:
            self.solver_exe = self._find_journal_runner()
        else:
            self.solver_exe = self._find_solver_executable()

        if not self.solver_exe.exists():
            raise FileNotFoundError(
                f"NX solver/runner not found at: {self.solver_exe}\n"
                f"Please check NX installation at: {self.nx_install_dir}"
            )

    def _find_nx_installation(self) -> Path:
        """Auto-detect NX installation directory."""
        # Common installation paths
        possible_paths = [
            Path(f"C:/Program Files/Siemens/NX{self.nastran_version}"),
            Path(f"C:/Program Files/Siemens/Simcenter3D_{self.nastran_version}"),
            Path(f"C:/Program Files (x86)/Siemens/NX{self.nastran_version}"),
        ]

        for path in possible_paths:
            if path.exists():
                return path

        raise FileNotFoundError(
            f"Could not auto-detect NX {self.nastran_version} installation.\n"
            f"Checked: {[str(p) for p in possible_paths]}\n"
            f"Please specify nx_install_dir manually."
        )

    def _find_journal_runner(self) -> Path:
        """Find the NX journal runner executable."""
        # First check the provided nx_install_dir
        if self.nx_install_dir:
            direct_path = self.nx_install_dir / "NXBIN" / "run_journal.exe"
            if direct_path.exists():
                return direct_path

        # Fallback: check common installation paths
        possible_exes = [
            Path(f"C:/Program Files/Siemens/Simcenter3D_{self.nastran_version}/NXBIN/run_journal.exe"),
            Path(f"C:/Program Files/Siemens/NX{self.nastran_version}/NXBIN/run_journal.exe"),
            Path(f"C:/Program Files/Siemens/DesigncenterNX{self.nastran_version}/NXBIN/run_journal.exe"),
        ]

        for exe in possible_exes:
            if exe.exists():
                return exe

        # Return the direct path (will error in __init__ if doesn't exist)
        return self.nx_install_dir / "NXBIN" / "run_journal.exe" if self.nx_install_dir else possible_exes[0]

    def _find_solver_executable(self) -> Path:
        """Find the Nastran solver executable."""
        # Use NX Nastran (not Simcenter) - has different licensing
        # Priority: Use NX installation, not Simcenter
        possible_exes = [
            self.nx_install_dir / "NXNASTRAN" / "bin" / "nastran.exe",
            self.nx_install_dir / "NXNASTRAN" / "nastran.exe",
            self.nx_install_dir / "bin" / "nastran.exe",
        ]

        for exe in possible_exes:
            if exe.exists():
                return exe

        # If not found in NX, try Simcenter as fallback
        simcenter_paths = [
            Path(f"C:/Program Files/Siemens/Simcenter3D_{self.nastran_version}"),
        ]

        for simcenter_dir in simcenter_paths:
            if simcenter_dir.exists():
                solve_exe = simcenter_dir / "NXNASTRAN" / "bin" / "nastran.exe"
                if solve_exe.exists():
                    return solve_exe

        # Return first guess (will error in __init__ if doesn't exist)
        return possible_exes[0]

    def create_iteration_folder(
        self,
        iterations_base_dir: Path,
        iteration_number: Optional[int] = None,
        expression_updates: Optional[Dict[str, float]] = None
    ) -> Path:
        """
        Create a fresh iteration folder with copies of all model files.

        HEEDS-style approach: Each iteration gets its own complete copy of model files.
        This ensures clean state and avoids model corruption between iterations.

        Folder structure created:
            iterN/
            ├── *.prt, *.fem, *.afm, *.sim  (copied from master)
            ├── params.exp                   (generated from expression_updates)
            ├── *.op2, *.f06                 (generated by solver)
            └── results/                     (for processed outputs like CSVs, HTMLs)

        Args:
            iterations_base_dir: Base directory for iteration folders (e.g., study/2_iterations)
            iteration_number: Specific iteration number, or auto-increment if None
            expression_updates: Dict of expression name -> value for params.exp file

        Returns:
            Path to the created iteration folder (e.g., 2_iterations/iter1)
        """
        if not self.master_model_dir:
            raise ValueError("master_model_dir must be set to use iteration folders")

        if not self.master_model_dir.exists():
            raise FileNotFoundError(f"Master model directory not found: {self.master_model_dir}")

        # Auto-increment iteration number if not provided
        if iteration_number is None:
            self._iteration_counter += 1
            iteration_number = self._iteration_counter

        # Create iteration folder: iter1, iter2, etc. (lowercase per user spec)
        iter_folder = iterations_base_dir / f"iter{iteration_number}"

        # Clean up if folder exists from a previous failed run
        if iter_folder.exists():
            print(f"[NX SOLVER] Cleaning up existing iteration folder: {iter_folder}")
            try:
                shutil.rmtree(iter_folder)
            except Exception as e:
                print(f"[NX SOLVER] WARNING: Could not clean up {iter_folder}: {e}")

        # Create fresh folder and results subfolder
        iter_folder.mkdir(parents=True, exist_ok=True)
        results_folder = iter_folder / "results"
        results_folder.mkdir(exist_ok=True)

        # Copy all NX model files from master
        print(f"[NX SOLVER] Creating iteration folder: {iter_folder.name}")
        print(f"[NX SOLVER]   Copying from: {self.master_model_dir}")

        copied_files = []
        for ext in self.MODEL_EXTENSIONS:
            for src_file in self.master_model_dir.glob(f"*{ext}"):
                dst_file = iter_folder / src_file.name
                try:
                    shutil.copy2(src_file, dst_file)
                    copied_files.append(src_file.name)
                except Exception as e:
                    print(f"[NX SOLVER]   WARNING: Could not copy {src_file.name}: {e}")

        print(f"[NX SOLVER]   Copied {len(copied_files)} model files")

        # Generate params.exp file if expression updates provided
        if expression_updates:
            exp_file = iter_folder / "params.exp"
            self._write_expression_file(exp_file, expression_updates)
            print(f"[NX SOLVER]   Generated params.exp with {len(expression_updates)} expressions")

        print(f"[NX SOLVER]   Created results/ subfolder for processed outputs")

        return iter_folder

    def _write_expression_file(self, exp_path: Path, expressions: Dict[str, float]):
        """
        Write expressions to .exp file format for NX import.

        Format: [unit]name=value
        Example: [mm]whiffle_min=42.5
        """
        # Default unit mapping (could be extended or made configurable)
        UNIT_MAPPING = {
            # Length parameters (mm)
            'whiffle_min': 'mm',
            'whiffle_triangle_closeness': 'mm',
            'inner_circular_rib_dia': 'mm',
            'outer_circular_rib_offset_from_outer': 'mm',
            'lateral_outer_pivot': 'mm',
            'lateral_inner_pivot': 'mm',
            'lateral_middle_pivot': 'mm',
            'lateral_closeness': 'mm',
            # Angle parameters (degrees)
            'whiffle_outer_to_vertical': 'Degrees',
            'lateral_inner_angle': 'Degrees',
            'lateral_outer_angle': 'Degrees',
            'blank_backface_angle': 'Degrees',
        }

        with open(exp_path, 'w') as f:
            for name, value in expressions.items():
                # Get unit, default to mm for unknown parameters
                unit = UNIT_MAPPING.get(name, 'mm')
                f.write(f"[{unit}]{name}={value}\n")

    def cleanup_iteration_folder(
        self,
        iter_folder: Path,
        keep_results: bool = True
    ):
        """
        Clean up an iteration folder after extracting results.

        Args:
            iter_folder: Path to iteration folder
            keep_results: If True, keeps .op2 and .f06 files; if False, deletes everything
        """
        if not iter_folder.exists():
            return

        if keep_results:
            # Delete everything except result files
            keep_extensions = {'.op2', '.f06', '.log'}
            for file in iter_folder.iterdir():
                if file.is_file() and file.suffix.lower() not in keep_extensions:
                    try:
                        file.unlink()
                    except Exception:
                        pass
            print(f"[NX SOLVER] Cleaned iteration folder (kept results): {iter_folder.name}")
        else:
            # Delete entire folder
            try:
                shutil.rmtree(iter_folder)
                print(f"[NX SOLVER] Deleted iteration folder: {iter_folder.name}")
            except Exception as e:
                print(f"[NX SOLVER] WARNING: Could not delete {iter_folder}: {e}")

    def run_simulation(
        self,
        sim_file: Path,
        working_dir: Optional[Path] = None,
        cleanup: bool = True,
        expression_updates: Optional[Dict[str, float]] = None,
        solution_name: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        Run NX Nastran simulation.

        Args:
            sim_file: Path to .sim file
            working_dir: Working directory for solver (defaults to sim file dir)
            cleanup: Remove intermediate files after solving
            expression_updates: Dict of expression name -> value to update
                              (only used in journal mode)
                              e.g., {'tip_thickness': 22.5, 'support_angle': 35.0}
            solution_name: Specific solution to solve (e.g., "Solution_Normal_Modes")
                          If None, solves all solutions. Only used in journal mode.

        Returns:
            Dictionary with:
                - success: bool
                - op2_file: Path to output .op2 file
                - log_file: Path to .log file
                - elapsed_time: Solve time in seconds
                - errors: List of error messages (if any)
                - solution_name: Name of the solution that was solved
        """
        sim_file = Path(sim_file)
        if not sim_file.exists():
            raise FileNotFoundError(f"Simulation file not found: {sim_file}")

        if working_dir is None:
            working_dir = sim_file.parent
        else:
            working_dir = Path(working_dir)

        # Check if we need to find/use .dat file (only in direct mode, not journal mode)
        # .sim files require NX GUI, but .dat files can be run directly with Nastran
        dat_file = None
        if not self.use_journal and sim_file.suffix == '.sim':
            # Look for corresponding .dat file (created by NX when solving)
            # Pattern: Bracket_sim1.sim -> bracket_sim1-solution_1.dat
            base = sim_file.stem.lower()
            possible_dats = list(working_dir.glob(f"{base}-solution_*.dat"))
            if possible_dats:
                # Use the most recent .dat file
                dat_file = max(possible_dats, key=lambda p: p.stat().st_mtime)
                print(f"\n[NX SOLVER] Found .dat file: {dat_file.name}")
                print(f"  Using .dat instead of .sim for better compatibility")
                sim_file = dat_file

        # Prepare output file names
        # When using journal mode with .sim files, output is named: <base>-solution_name.op2
        # When using direct mode with .dat files, output is named: <base>.op2
        base_name = sim_file.stem

        if self.use_journal and sim_file.suffix == '.sim':
            # Journal mode: determine solution-specific output name
            if solution_name:
                # Convert solution name to lowercase and replace spaces with underscores
                # E.g., "Solution_Normal_Modes" -> "solution_normal_modes"
                solution_suffix = solution_name.lower().replace(' ', '_')
                output_base = f"{base_name.lower()}-{solution_suffix}"
            else:
                # Default to solution_1
                output_base = f"{base_name.lower()}-solution_1"
        else:
            # Direct mode or .dat file
            output_base = base_name

        op2_file = working_dir / f"{output_base}.op2"
        log_file = working_dir / f"{output_base}.log"
        f06_file = working_dir / f"{output_base}.f06"

        print(f"\n[NX SOLVER] Starting simulation...")
        print(f"  Input file: {sim_file.name}")
        print(f"  Working dir: {working_dir}")
        print(f"  Mode: {'Journal' if self.use_journal else 'Direct'}")

        # Record timestamps of old files BEFORE solving
        # We'll verify files are regenerated by checking timestamps AFTER solve
        # This is more reliable than deleting (which can fail due to file locking on Windows)
        old_op2_time = op2_file.stat().st_mtime if op2_file.exists() else None
        old_f06_time = f06_file.stat().st_mtime if f06_file.exists() else None
        old_log_time = log_file.stat().st_mtime if log_file.exists() else None

        if old_op2_time:
            print(f"  Found existing OP2 (modified: {time.ctime(old_op2_time)})")
            print(f"  Will verify NX regenerates it with newer timestamp")

        # Build command based on mode
        if self.use_journal and sim_file.suffix == '.sim':
            # Use NX journal for .sim files (handles licensing properly)
            # Generate a temporary journal file with the correct sim file path
            journal_template = Path(__file__).parent / "solve_simulation.py"
            temp_journal = working_dir / "_temp_solve_journal.py"

            # Read template and replace placeholder with actual path
            with open(journal_template, 'r') as f:
                journal_content = f.read()

            # Create a custom journal that passes the sim file path, solution name, and expression values
            # Build argv list with expression updates
            argv_list = [f"r'{sim_file.absolute()}'"]

            # Add solution name if provided (passed as second argument)
            if solution_name:
                argv_list.append(f"'{solution_name}'")
            else:
                argv_list.append("None")

            # Add expression values if provided
            # Pass all expressions as key=value pairs
            if expression_updates:
                for expr_name, expr_value in expression_updates.items():
                    argv_list.append(f"'{expr_name}={expr_value}'")

            argv_str = ', '.join(argv_list)

            custom_journal = f'''# Auto-generated journal for solving {sim_file.name}
import sys
sys.argv = ['', {argv_str}]  # Set argv for the main function
{journal_content}
'''
            with open(temp_journal, 'w') as f:
                f.write(custom_journal)

            cmd = [
                str(self.solver_exe),  # run_journal.exe
                str(temp_journal.absolute())  # Use absolute path to avoid path issues
            ]
        else:
            # Direct Nastran batch command for .dat files or direct mode
            # IMPORTANT: prog=bundle enables bundle licensing (required for desktop licenses)
            cmd = [
                str(self.solver_exe),
                str(sim_file),
                "prog=bundle",
                "old=no",
                "scratch=yes"
            ]

        # Set up environment for Simcenter/NX
        env = os.environ.copy()

        # Get SPLM_LICENSE_SERVER - prefer system registry (most up-to-date) over process env
        license_server = ''

        # First try system-level environment (Windows registry) - this is the authoritative source
        import subprocess as sp
        try:
            result = sp.run(
                ['reg', 'query', 'HKLM\\SYSTEM\\CurrentControlSet\\Control\\Session Manager\\Environment', '/v', 'SPLM_LICENSE_SERVER'],
                capture_output=True, text=True, timeout=5
            )
            if result.returncode == 0:
                # Parse: "    SPLM_LICENSE_SERVER    REG_SZ    value"
                for line in result.stdout.splitlines():
                    if 'SPLM_LICENSE_SERVER' in line:
                        parts = line.split('REG_SZ')
                        if len(parts) > 1:
                            license_server = parts[1].strip()
                            break
        except Exception:
            pass

        # Fall back to process environment if registry query failed
        if not license_server:
            license_server = env.get('SPLM_LICENSE_SERVER', '')

        if license_server:
            env['SPLM_LICENSE_SERVER'] = license_server
            print(f"[NX SOLVER] Using license server: {license_server}")
        else:
            env['SPLM_LICENSE_SERVER'] = '29000@localhost'
            print(f"[NX SOLVER] WARNING: SPLM_LICENSE_SERVER not set, using default: {env['SPLM_LICENSE_SERVER']}")

        # Add NX/Simcenter paths to environment
        nx_bin = self.nx_install_dir / "NXBIN"
        if nx_bin.exists():
            env['PATH'] = f"{nx_bin};{env.get('PATH', '')}"

        nastran_bin = self.solver_exe.parent
        if nastran_bin.exists():
            env['PATH'] = f"{nastran_bin};{env.get('PATH', '')}"

        # Run solver
        start_time = time.time()

        try:
            result = subprocess.run(
                cmd,
                cwd=str(working_dir),
                capture_output=True,
                text=True,
                timeout=self.timeout,
                env=env  # Use modified environment
            )

            elapsed_time = time.time() - start_time

            # Display journal output for debugging
            if self.use_journal:
                if result.stdout and result.stdout.strip():
                    print("[JOURNAL OUTPUT]")
                    for line in result.stdout.strip().split('\n'):
                        print(f"  {line}")

            # Check for journal errors
            if self.use_journal and result.stderr and "error" in result.stderr.lower():
                print("[JOURNAL ERRORS]")
                for line in result.stderr.strip().split('\n')[:5]:
                    print(f"  {line}")

            # Wait for output files to appear AND be regenerated (journal mode runs solve in background)
            if self.use_journal:
                max_wait = 30  # seconds - background solves can take time
                wait_start = time.time()
                print("[NX SOLVER] Waiting for solve to complete...")

                # Wait for files to exist AND have newer timestamps than before
                while (time.time() - wait_start) < max_wait:
                    files_exist = f06_file.exists() and op2_file.exists()

                    if files_exist:
                        # Verify files were regenerated (newer timestamps)
                        new_op2_time = op2_file.stat().st_mtime
                        new_f06_time = f06_file.stat().st_mtime

                        # If no old files, or new files are newer, we're done!
                        if (old_op2_time is None or new_op2_time > old_op2_time) and \
                           (old_f06_time is None or new_f06_time > old_f06_time):
                            elapsed = time.time() - wait_start
                            print(f"[NX SOLVER] Fresh output files detected after {elapsed:.1f}s")
                            if old_op2_time:
                                print(f"  OP2 regenerated: {time.ctime(old_op2_time)} -> {time.ctime(new_op2_time)}")
                            break

                    time.sleep(0.5)
                    if (time.time() - wait_start) % 2 < 0.5:  # Print every 2 seconds
                        elapsed = time.time() - wait_start
                        print(f"  Waiting for fresh results... ({elapsed:.0f}s)")

                # Final check - FAIL if files weren't regenerated
                op2_is_fresh = True
                if op2_file.exists():
                    current_op2_time = op2_file.stat().st_mtime
                    if old_op2_time and current_op2_time <= old_op2_time:
                        print(f"  ERROR: OP2 file was NOT regenerated! (Still has old timestamp)")
                        print(f"         Old: {time.ctime(old_op2_time)}, Current: {time.ctime(current_op2_time)}")
                        print(f"         The solve failed - cannot use stale results!")
                        op2_is_fresh = False
                else:
                    print(f"  ERROR: OP2 file does not exist!")
                    op2_is_fresh = False

            # Check for completion - also require fresh OP2
            success = self._check_solution_success(f06_file, log_file) and op2_is_fresh

            errors = []
            if not success:
                errors = self._extract_errors(f06_file, log_file)

            # Clean up intermediate files if requested
            if cleanup and success:
                self._cleanup_temp_files(working_dir, base_name)

            # Clean up temporary journal file if it was created
            temp_journal_path = working_dir / "_temp_solve_journal.py"
            if temp_journal_path.exists():
                try:
                    temp_journal_path.unlink()
                except Exception:
                    pass

            print(f"[NX SOLVER] Complete in {elapsed_time:.1f}s")
            if success:
                print(f"[NX SOLVER] Results: {op2_file.name}")
            else:
                print(f"[NX SOLVER] FAILED - check {f06_file.name}")
                for error in errors:
                    print(f"  ERROR: {error}")

            return {
                'success': success,
                'op2_file': op2_file if op2_file.exists() else None,
                'log_file': log_file if log_file.exists() else None,
                'f06_file': f06_file if f06_file.exists() else None,
                'elapsed_time': elapsed_time,
                'errors': errors,
                'return_code': result.returncode,
                'solution_name': solution_name
            }

        except subprocess.TimeoutExpired:
            elapsed_time = time.time() - start_time
            print(f"[NX SOLVER] TIMEOUT after {elapsed_time:.1f}s")
            return {
                'success': False,
                'op2_file': None,
                'log_file': log_file if log_file.exists() else None,
                'elapsed_time': elapsed_time,
                'errors': [f'Solver timeout after {self.timeout}s'],
                'return_code': -1,
                'solution_name': solution_name
            }

        except Exception as e:
            elapsed_time = time.time() - start_time
            print(f"[NX SOLVER] ERROR: {e}")
            return {
                'success': False,
                'op2_file': None,
                'log_file': None,
                'elapsed_time': elapsed_time,
                'errors': [str(e)],
                'return_code': -1,
                'solution_name': solution_name
            }

    def _check_solution_success(self, f06_file: Path, log_file: Path) -> bool:
        """
        Check if solution completed successfully.

        Looks for completion markers in .f06 and .log files.
        """
        # Check .f06 file for completion
        if f06_file.exists():
            try:
                with open(f06_file, 'r', encoding='latin-1', errors='ignore') as f:
                    content = f.read()
                    # Look for successful completion markers
                    if 'NORMAL TERMINATION' in content or 'USER INFORMATION MESSAGE' in content:
                        return True
                    # Check for fatal errors
                    if 'FATAL MESSAGE' in content or 'EXECUTION TERMINATED' in content:
                        return False
            except Exception:
                pass

        # Fallback: check if OP2 was created recently
        op2_file = f06_file.with_suffix('.op2')
        if op2_file.exists():
            # If OP2 was modified within last minute, assume success
            if (time.time() - op2_file.stat().st_mtime) < 60:
                return True

        return False

    def _extract_errors(self, f06_file: Path, log_file: Path) -> list:
        """Extract error messages from output files."""
        errors = []

        if f06_file.exists():
            try:
                with open(f06_file, 'r', encoding='latin-1', errors='ignore') as f:
                    for line in f:
                        if 'FATAL' in line or 'ERROR' in line:
                            errors.append(line.strip())
            except Exception:
                pass

        return errors[:10]  # Limit to first 10 errors

    def discover_model(self, sim_file: Path) -> Dict[str, Any]:
        """
        Discover model information without solving.

        This scans the NX simulation file and reports:
        - All solutions (names, types)
        - All expressions (potential design variables)
        - Mesh info
        - Linked geometry parts

        Args:
            sim_file: Path to .sim file

        Returns:
            Dictionary with discovered model info
        """
        import json

        sim_file = Path(sim_file)
        if not sim_file.exists():
            return {'success': False, 'error': f'Sim file not found: {sim_file}'}

        # Use the discover_model journal
        discover_journal = Path(__file__).parent.parent / "nx_journals" / "discover_model.py"

        if not discover_journal.exists():
            return {'success': False, 'error': f'Discovery journal not found: {discover_journal}'}

        print(f"\n[NX SOLVER] Discovering model: {sim_file.name}")
        print(f"  Using journal: {discover_journal.name}")

        try:
            cmd = [str(self.solver_exe), str(discover_journal), '--', str(sim_file.absolute())]

            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=60,  # 1 minute timeout for discovery
                cwd=str(sim_file.parent)
            )

            # Print stderr (debug/progress messages)
            if result.stderr:
                for line in result.stderr.strip().split('\n'):
                    print(f"  {line}")

            # Parse stdout as JSON
            if result.stdout:
                try:
                    discovery_result = json.loads(result.stdout)
                    return discovery_result
                except json.JSONDecodeError as e:
                    return {
                        'success': False,
                        'error': f'Failed to parse discovery output: {e}',
                        'raw_output': result.stdout[:1000]
                    }
            else:
                return {
                    'success': False,
                    'error': 'No output from discovery journal',
                    'stderr': result.stderr
                }

        except subprocess.TimeoutExpired:
            return {'success': False, 'error': 'Discovery timeout (60s)'}
        except Exception as e:
            return {'success': False, 'error': str(e)}

    def _cleanup_temp_files(self, working_dir: Path, base_name: str):
        """Remove temporary solver files."""
        # Files to keep
        keep_extensions = {'.op2', '.f06', '.log'}

        # Files to remove
        remove_patterns = [
            f"{base_name}.f04",
            f"{base_name}.dat",
            f"{base_name}.diag",
            f"{base_name}.master",
            f"{base_name}.dball",
            f"{base_name}.MASTER",
            f"{base_name}.DBALL",
            f"{base_name}_*.png",
            f"{base_name}_*.html",
        ]

        for pattern in remove_patterns:
            for file in working_dir.glob(pattern):
                try:
                    file.unlink()
                except Exception:
                    pass


# Convenience function for optimization loops
def run_nx_simulation(
    sim_file: Path,
    nastran_version: str = "2412",
    timeout: int = 600,
    cleanup: bool = True,
    use_journal: bool = True,
    expression_updates: Optional[Dict[str, float]] = None,
    solution_name: Optional[str] = None
) -> Path:
    """
    Convenience function to run NX simulation and return OP2 file path.

    Args:
        sim_file: Path to .sim file
        nastran_version: NX version
        timeout: Solver timeout in seconds
        cleanup: Remove temp files
        use_journal: Use NX journal for solving (recommended for licensing)
        expression_updates: Dict of expression name -> value to update in journal
        solution_name: Specific solution to solve (e.g., "Solution_Normal_Modes")

    Returns:
        Path to output .op2 file

    Raises:
        RuntimeError: If simulation fails
    """
    solver = NXSolver(nastran_version=nastran_version, timeout=timeout, use_journal=use_journal)
    result = solver.run_simulation(
        sim_file,
        cleanup=cleanup,
        expression_updates=expression_updates,
        solution_name=solution_name
    )

    if not result['success']:
        error_msg = '\n'.join(result['errors']) if result['errors'] else 'Unknown error'
        raise RuntimeError(f"NX simulation failed:\n{error_msg}")

    if not result['op2_file'] or not result['op2_file'].exists():
        raise RuntimeError("Simulation completed but OP2 file not found")

    return result['op2_file']