Atomizer/optimization_engine/nx_solver.py

"""
NX Nastran Solver Integration

Executes NX Nastran solver in batch mode for optimization loops.
"""

from pathlib import Path
from typing import Optional, Dict, Any
import subprocess
import time
import shutil
import os


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

    def __init__(
        self,
        nx_install_dir: Optional[Path] = None,
        nastran_version: str = "2412",
        timeout: int = 600,
        use_journal: bool = True
    ):
        """
        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)
        """
        self.nastran_version = nastran_version
        self.timeout = timeout
        self.use_journal = use_journal

        # 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."""
        # Simcenter3D has run_journal.exe for batch execution
        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"),
        ]

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

        # Return first guess (will error in __init__ if doesn't exist)
        return 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 run_simulation(
        self,
        sim_file: Path,
        working_dir: Optional[Path] = None,
        cleanup: bool = True,
        expression_updates: Optional[Dict[str, float]] = 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}

        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)
        """
        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_1.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: look for -solution_1 pattern
            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'}")

        # Delete old result files (.op2, .log, .f06) to force fresh solve
        # (.dat file is needed by NX, don't delete it!)
        # (Otherwise NX may reuse cached results)
        files_to_delete = [op2_file, log_file, f06_file]

        deleted_count = 0
        for old_file in files_to_delete:
            if old_file.exists():
                try:
                    old_file.unlink()
                    deleted_count += 1
                except Exception as e:
                    print(f"  Warning: Could not delete {old_file.name}: {e}")

        if deleted_count > 0:
            print(f"  Deleted {deleted_count} old result file(s) to force fresh solve")

        # 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 and expression values
            # Build argv list with expression updates
            argv_list = [f"r'{sim_file.absolute()}'"]

            # Add expression values if provided
            if expression_updates:
                # For bracket example, we expect: tip_thickness, support_angle
                if 'tip_thickness' in expression_updates:
                    argv_list.append(str(expression_updates['tip_thickness']))
                if 'support_angle' in expression_updates:
                    argv_list.append(str(expression_updates['support_angle']))

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

        # Set license server (use 29000 for Simcenter)
        # Override any incorrect license server settings
        env['SPLM_LICENSE_SERVER'] = '29000@AntoineThinkpad'

        # Force desktop licensing instead of enterprise
        # User has nx_nas_bn_basic_dsk (desktop) not nx_nas_basic_ent (enterprise)
        env['NXNA_LICENSE_FILE'] = '29000@AntoineThinkpad'
        env['NXNASTRAN_LICENSE_FILE'] = '29000@AntoineThinkpad'

        # 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 (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...")
                while not (f06_file.exists() and op2_file.exists()) and (time.time() - wait_start) < max_wait:
                    time.sleep(0.5)
                    if (time.time() - wait_start) % 2 < 0.5:  # Print every 2 seconds
                        elapsed = time.time() - wait_start
                        print(f"  Waiting... ({elapsed:.0f}s)")

                if f06_file.exists() and op2_file.exists():
                    print(f"[NX SOLVER] Output files detected after {time.time() - wait_start:.1f}s")

            # Check for completion
            success = self._check_solution_success(f06_file, log_file)

            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
            }

        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
            }

        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
            }

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

    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)

    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']