Atomizer/optimization_engine/hooks/nx_cae/solver_manager.py

"""
NX Solver Manager Hook
======================

Provides Python functions to export BDF decks and solve simulations.

API Reference (NX Open):
- SimSolution.ExportSolver() -> Export Nastran deck (.dat/.bdf)
- SimSolution.Solve() -> Solve a single solution
- SimSolveManager.SolveChainOfSolutions() -> Solve solution chain

Phase 2 Task 2.1 - NX Open Automation Roadmap

Usage:
    from optimization_engine.hooks.nx_cae import solver_manager

    # Export BDF without solving
    result = solver_manager.export_bdf(
        "C:/model.sim",
        "Solution 1",
        "C:/output/model.dat"
    )

    # Solve simulation
    result = solver_manager.solve_simulation("C:/model.sim", "Solution 1")
"""

import os
import json
import subprocess
import tempfile
from pathlib import Path
from typing import Optional, Dict, Any

# NX installation path (configurable)
NX_BIN_PATH = os.environ.get(
    "NX_BIN_PATH",
    r"C:\Program Files\Siemens\NX2506\NXBIN"
)

# Journal template for BDF export
BDF_EXPORT_JOURNAL = '''
# NX Open Python Journal - BDF Export
# Auto-generated by Atomizer hooks
# Phase 2 Task 2.1 - NX Open Automation Roadmap

import NXOpen
import NXOpen.CAE
import json
import sys
import os


def main():
    """Export BDF/DAT file from a simulation solution."""
    args = sys.argv[1:] if len(sys.argv) > 1 else []

    if len(args) < 3:
        raise ValueError("Usage: script.py <sim_path> <solution_name> <output_bdf> [output_json]")

    sim_path = args[0]
    solution_name = args[1]
    output_bdf = args[2]
    output_json = args[3] if len(args) > 3 else None

    result = {"success": False, "error": None, "data": {}}

    try:
        session = NXOpen.Session.GetSession()

        # Set load options
        working_dir = os.path.dirname(sim_path)
        session.Parts.LoadOptions.ComponentLoadMethod = NXOpen.LoadOptions.LoadMethod.FromDirectory
        session.Parts.LoadOptions.SetSearchDirectories([working_dir], [True])

        # Open the simulation file
        print(f"[JOURNAL] Opening simulation: {sim_path}")
        basePart, loadStatus = session.Parts.OpenActiveDisplay(
            sim_path,
            NXOpen.DisplayPartOption.AllowAdditional
        )
        loadStatus.Dispose()

        # Get the sim part
        simPart = session.Parts.Work
        if not isinstance(simPart, NXOpen.CAE.SimPart):
            raise ValueError(f"Part is not a SimPart: {type(simPart)}")

        simSimulation = simPart.Simulation
        print(f"[JOURNAL] Simulation: {simSimulation.Name}")

        # Find the solution
        solution = None
        for sol in simSimulation.Solutions:
            if sol.Name == solution_name:
                solution = sol
                break

        if solution is None:
            # Try to find by index or use first solution
            solutions = list(simSimulation.Solutions)
            if solutions:
                solution = solutions[0]
                print(f"[JOURNAL] Solution '{solution_name}' not found, using '{solution.Name}'")
            else:
                raise ValueError(f"No solutions found in simulation")

        print(f"[JOURNAL] Solution: {solution.Name}")

        # Export the solver deck
        # The ExportSolver method exports the Nastran input deck
        print(f"[JOURNAL] Exporting BDF to: {output_bdf}")

        # Create export builder
        # NX API: SimSolution has methods for exporting
        # Method 1: Try ExportSolver if available
        try:
            # Some NX versions use NastranSolverExportBuilder
            exportBuilder = solution.CreateNastranSolverExportBuilder()
            exportBuilder.NastranInputFile = output_bdf
            exportBuilder.Commit()
            exportBuilder.Destroy()
            print("[JOURNAL] Exported via NastranSolverExportBuilder")
        except AttributeError:
            # Method 2: Alternative - solve and copy output
            # When solving, NX creates the deck in SXXXXX folder
            print("[JOURNAL] NastranSolverExportBuilder not available")
            print("[JOURNAL] BDF export requires solving - use solve_simulation instead")
            raise ValueError("Direct BDF export not available in this NX version. "
                           "Use solve_simulation() and find BDF in solution folder.")

        result["success"] = True
        result["data"] = {
            "output_file": output_bdf,
            "solution_name": solution.Name,
            "simulation": simSimulation.Name,
        }
        print(f"[JOURNAL] Export completed successfully")

    except Exception as e:
        result["error"] = str(e)
        print(f"[JOURNAL] ERROR: {e}")
        import traceback
        traceback.print_exc()

    # Write result
    if output_json:
        with open(output_json, 'w') as f:
            json.dump(result, f, indent=2)

    return result


if __name__ == '__main__':
    main()
'''


def _run_journal(journal_content: str, *args) -> Dict[str, Any]:
    """Execute an NX journal script and return the result."""
    run_journal_exe = Path(NX_BIN_PATH) / "run_journal.exe"
    if not run_journal_exe.exists():
        return {
            "success": False,
            "error": f"run_journal.exe not found at {run_journal_exe}",
            "data": {}
        }

    # Create temporary files
    with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as journal_file:
        journal_file.write(journal_content)
        journal_path = journal_file.name

    with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as output_file:
        output_path = output_file.name

    try:
        # Build command
        cmd = [str(run_journal_exe), journal_path, "-args"]
        cmd.extend(str(a) for a in args)
        cmd.append(output_path)

        # Execute
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=120  # 2 minute timeout
        )

        # Read result
        if os.path.exists(output_path):
            with open(output_path, 'r') as f:
                return json.load(f)
        else:
            return {
                "success": False,
                "error": f"No output file generated. stdout: {result.stdout}, stderr: {result.stderr}",
                "data": {}
            }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "error": "Journal execution timed out after 120 seconds",
            "data": {}
        }
    except Exception as e:
        return {
            "success": False,
            "error": str(e),
            "data": {}
        }
    finally:
        # Cleanup
        try:
            os.unlink(journal_path)
        except:
            pass
        try:
            os.unlink(output_path)
        except:
            pass


def export_bdf(
    sim_path: str,
    solution_name: str = "Solution 1",
    output_bdf: Optional[str] = None
) -> Dict[str, Any]:
    """
    Export Nastran deck (BDF/DAT) from a simulation without solving.

    Note: This functionality depends on NX version. Some versions require
    solving to generate the BDF. Use solve_simulation() and locate the BDF
    in the solution folder (SXXXXX/*.dat) as an alternative.

    Args:
        sim_path: Path to .sim file
        solution_name: Name of solution to export (default "Solution 1")
        output_bdf: Output path for BDF file (default: same dir as sim)

    Returns:
        dict: {
            'success': bool,
            'error': str or None,
            'data': {
                'output_file': Path to exported BDF,
                'solution_name': Solution name used,
                'simulation': Simulation name
            }
        }

    Example:
        >>> result = export_bdf("C:/model.sim", "Solution 1", "C:/output/model.dat")
        >>> if result["success"]:
        ...     print(f"BDF exported to: {result['data']['output_file']}")
    """
    sim_path = str(Path(sim_path).resolve())

    if not Path(sim_path).exists():
        return {
            "success": False,
            "error": f"Simulation file not found: {sim_path}",
            "data": {}
        }

    if output_bdf is None:
        sim_dir = Path(sim_path).parent
        sim_name = Path(sim_path).stem
        output_bdf = str(sim_dir / f"{sim_name}.dat")

    return _run_journal(BDF_EXPORT_JOURNAL, sim_path, solution_name, output_bdf)


def get_bdf_from_solution_folder(
    sim_path: str,
    solution_name: str = "Solution 1"
) -> Dict[str, Any]:
    """
    Locate BDF file in the solution output folder.

    After solving, NX creates a folder structure like:
    - model_sim1_fem1_SXXXXX/
      - model_sim1_fem1.dat  (BDF file)
      - model_sim1_fem1.op2  (results)

    This function finds the BDF without running export.

    Args:
        sim_path: Path to .sim file
        solution_name: Name of solution

    Returns:
        dict: {
            'success': bool,
            'error': str or None,
            'data': {
                'bdf_file': Path to BDF if found,
                'solution_folders': List of found solution folders
            }
        }
    """
    sim_path = Path(sim_path)
    if not sim_path.exists():
        return {
            "success": False,
            "error": f"Simulation file not found: {sim_path}",
            "data": {}
        }

    sim_dir = sim_path.parent
    sim_stem = sim_path.stem

    # Search for solution folders (pattern: *_SXXXXX)
    solution_folders = list(sim_dir.glob(f"{sim_stem}*_S[0-9]*"))

    if not solution_folders:
        # Also try simpler patterns
        solution_folders = list(sim_dir.glob("*_S[0-9]*"))

    bdf_files = []
    for folder in solution_folders:
        if folder.is_dir():
            # Look for .dat or .bdf files
            dat_files = list(folder.glob("*.dat"))
            bdf_files.extend(dat_files)

    if bdf_files:
        # Return the most recent one
        bdf_files.sort(key=lambda f: f.stat().st_mtime, reverse=True)
        return {
            "success": True,
            "error": None,
            "data": {
                "bdf_file": str(bdf_files[0]),
                "all_bdf_files": [str(f) for f in bdf_files],
                "solution_folders": [str(f) for f in solution_folders]
            }
        }
    else:
        return {
            "success": False,
            "error": "No BDF files found. Ensure the simulation has been solved.",
            "data": {
                "solution_folders": [str(f) for f in solution_folders]
            }
        }


def solve_simulation(
    sim_path: str,
    solution_name: str = "Solution 1",
    expression_updates: Optional[Dict[str, float]] = None
) -> Dict[str, Any]:
    """
    Solve a simulation solution.

    This uses the existing solve_simulation.py journal which handles both
    single-part and assembly FEM workflows.

    Args:
        sim_path: Path to .sim file
        solution_name: Name of solution to solve (default "Solution 1")
        expression_updates: Optional dict of {expression_name: value} to update

    Returns:
        dict: {
            'success': bool,
            'error': str or None,
            'data': {
                'solution_folder': Path to solution output folder,
                'op2_file': Path to OP2 results file,
                'bdf_file': Path to BDF input file
            }
        }

    Note:
        For full solve functionality, use the NXSolver class in
        optimization_engine/nx_solver.py which provides more features
        like iteration folders and batch processing.
    """
    # This is a simplified wrapper - for full functionality use NXSolver
    solve_journal = Path(__file__).parent.parent.parent / "solve_simulation.py"

    if not solve_journal.exists():
        return {
            "success": False,
            "error": f"Solve journal not found: {solve_journal}",
            "data": {}
        }

    run_journal_exe = Path(NX_BIN_PATH) / "run_journal.exe"
    if not run_journal_exe.exists():
        return {
            "success": False,
            "error": f"run_journal.exe not found at {run_journal_exe}",
            "data": {}
        }

    # Build command
    cmd = [str(run_journal_exe), str(solve_journal), "-args", sim_path, solution_name]

    # Add expression updates
    if expression_updates:
        for name, value in expression_updates.items():
            cmd.append(f"{name}={value}")

    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=600  # 10 minute timeout for solving
        )

        # Check for success in output
        if "Solve completed successfully" in result.stdout or result.returncode == 0:
            # Find output files
            bdf_result = get_bdf_from_solution_folder(sim_path, solution_name)

            return {
                "success": True,
                "error": None,
                "data": {
                    "stdout": result.stdout[-2000:],  # Last 2000 chars
                    "bdf_file": bdf_result["data"].get("bdf_file") if bdf_result["success"] else None,
                    "solution_folders": bdf_result["data"].get("solution_folders", [])
                }
            }
        else:
            return {
                "success": False,
                "error": f"Solve may have failed. Check output.",
                "data": {
                    "stdout": result.stdout[-2000:],
                    "stderr": result.stderr[-1000:]
                }
            }

    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "error": "Solve timed out after 600 seconds",
            "data": {}
        }
    except Exception as e:
        return {
            "success": False,
            "error": str(e),
            "data": {}
        }


if __name__ == "__main__":
    # Example usage
    import sys
    if len(sys.argv) > 1:
        sim_path = sys.argv[1]
        solution = sys.argv[2] if len(sys.argv) > 2 else "Solution 1"

        print(f"Looking for BDF in solution folder...")
        result = get_bdf_from_solution_folder(sim_path, solution)

        if result["success"]:
            print(f"Found BDF: {result['data']['bdf_file']}")
        else:
            print(f"Error: {result['error']}")
            print(f"Trying to export...")
            result = export_bdf(sim_path, solution)
            print(f"Export result: {result}")
    else:
        print("Usage: python solver_manager.py <sim_path> [solution_name]")