Atomizer/optimization_engine/hooks/nx_cad/geometry_query.py

"""
NX Geometry Query Hook
======================

Provides Python functions to query geometry properties from NX parts.

API Reference (verified via Siemens MCP docs):
- Part.MeasureManager() -> Returns measure manager for this part
- MeasureManager.NewMassProperties() -> Create mass properties measurement
- Part.Bodies() -> BodyCollection (solid bodies in the part)
- Body.GetPhysicalMaterial() -> Get material assigned to body

Usage:
    from optimization_engine.hooks.nx_cad import geometry_query

    # Get mass properties
    result = geometry_query.get_mass_properties("C:/path/to/part.prt")

    # Get body info
    result = geometry_query.get_bodies("C:/path/to/part.prt")

    # Get volume
    result = geometry_query.get_volume("C:/path/to/part.prt")
"""

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

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

# Journal template for geometry query operations
GEOMETRY_QUERY_JOURNAL = '''
# NX Open Python Journal - Geometry Query Operations
# Auto-generated by Atomizer hooks
#
# Based on Siemens NX Open Python API:
# - MeasureManager.NewMassProperties()
# - BodyCollection
# - Body.GetPhysicalMaterial()

import NXOpen
import NXOpen.UF
import json
import sys
import os
import math

def main():
    """Execute geometry query operation based on command arguments."""
    # Get the NX session
    session = NXOpen.Session.GetSession()

    # Parse arguments: operation, part_path, output_json, [extra_args...]
    args = sys.argv[1:] if len(sys.argv) > 1 else []

    if len(args) < 3:
        raise ValueError("Usage: script.py <operation> <part_path> <output_json> [args...]")

    operation = args[0]
    part_path = args[1]
    output_json = args[2]
    extra_args = args[3:] if len(args) > 3 else []

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

    try:
        # Ensure part is open
        part = ensure_part_open(session, part_path)

        if part is None:
            result["error"] = f"Failed to open part: {part_path}"
        elif operation == "mass_properties":
            result = get_mass_properties(part)
        elif operation == "bodies":
            result = get_bodies(part)
        elif operation == "volume":
            result = get_volume(part)
        elif operation == "surface_area":
            result = get_surface_area(part)
        elif operation == "material":
            result = get_material(part)
        else:
            result["error"] = f"Unknown operation: {operation}"

    except Exception as e:
        import traceback
        result["error"] = str(e)
        result["traceback"] = traceback.format_exc()

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

    return result


def ensure_part_open(session, part_path):
    """Ensure the part is open and return it."""
    # Check if already open
    part_path_normalized = os.path.normpath(part_path).lower()

    for part in session.Parts:
        if os.path.normpath(part.FullPath).lower() == part_path_normalized:
            return part

    # Need to open it
    if not os.path.exists(part_path):
        return None

    try:
        # Set load options for the working directory
        working_dir = os.path.dirname(part_path)
        session.Parts.LoadOptions.ComponentLoadMethod = NXOpen.LoadOptions.LoadMethod.FromDirectory
        session.Parts.LoadOptions.SetSearchDirectories([working_dir], [True])

        # Use OpenActiveDisplay instead of OpenBase for better compatibility
        part, load_status = session.Parts.OpenActiveDisplay(
            part_path,
            NXOpen.DisplayPartOption.AllowAdditional
        )
        load_status.Dispose()
        return part
    except:
        return None


def get_solid_bodies(part):
    """Get all solid bodies from a part."""
    solid_bodies = []
    for body in part.Bodies:
        if body.IsSolidBody:
            solid_bodies.append(body)
    return solid_bodies


def get_mass_properties(part):
    """Get mass properties from a part.

    NX Open API: MeasureManager.NewMassProperties()

    Returns mass, volume, surface area, centroid, and inertia properties.
    """
    result = {"success": False, "error": None, "data": {}}

    try:
        # Get solid bodies
        solid_bodies = get_solid_bodies(part)

        if not solid_bodies:
            result["error"] = "No solid bodies found in part"
            return result

        # Get measure manager
        measure_manager = part.MeasureManager

        # Get units - use base units array like the working journal
        uc = part.UnitCollection
        mass_units = [
            uc.GetBase("Area"),
            uc.GetBase("Volume"),
            uc.GetBase("Mass"),
            uc.GetBase("Length")
        ]

        # Create mass properties measurement
        # Signature: NewMassProperties(mass_units, accuracy, objects)
        mass_props = measure_manager.NewMassProperties(mass_units, 0.99, solid_bodies)

        # Get properties
        mass = mass_props.Mass
        volume = mass_props.Volume
        area = mass_props.Area

        # Get centroid
        centroid = mass_props.Centroid
        centroid_x = centroid.X
        centroid_y = centroid.Y
        centroid_z = centroid.Z

        # Get principal moments of inertia (may not be available)
        ixx = 0.0
        iyy = 0.0
        izz = 0.0
        try:
            principal_moments = mass_props.PrincipalMomentsOfInertia
            ixx = principal_moments[0]
            iyy = principal_moments[1]
            izz = principal_moments[2]
        except:
            pass

        # Get material info from first body via attributes
        material_name = None
        density = None
        try:
            # Try body attributes (NX stores material as attribute)
            attrs = solid_bodies[0].GetUserAttributes()
            for attr in attrs:
                if 'material' in attr.Title.lower():
                    material_name = attr.StringValue
                    break
        except:
            pass

        result["success"] = True
        result["data"] = {
            "mass": mass,
            "mass_unit": "kg",
            "volume": volume,
            "volume_unit": "mm^3",
            "surface_area": area,
            "area_unit": "mm^2",
            "centroid": {
                "x": centroid_x,
                "y": centroid_y,
                "z": centroid_z,
                "unit": "mm"
            },
            "principal_moments": {
                "Ixx": ixx,
                "Iyy": iyy,
                "Izz": izz,
                "unit": "kg*mm^2"
            },
            "material": material_name,
            "density": density,
            "body_count": len(solid_bodies)
        }

    except Exception as e:
        import traceback
        result["error"] = str(e)
        result["traceback"] = traceback.format_exc()

    return result


def get_bodies(part):
    """Get information about all bodies in the part.

    NX Open API: Part.Bodies()
    """
    result = {"success": False, "error": None, "data": {}}

    try:
        bodies_info = []

        for body in part.Bodies:
            body_data = {
                "name": body.Name if hasattr(body, 'Name') else None,
                "is_solid": body.IsSolidBody,
                "is_sheet": body.IsSheetBody,
            }

            # Try to get material
            try:
                phys_mat = body.GetPhysicalMaterial()
                if phys_mat:
                    body_data["material"] = phys_mat.Name
            except:
                body_data["material"] = None

            bodies_info.append(body_data)

        result["success"] = True
        result["data"] = {
            "count": len(bodies_info),
            "solid_count": sum(1 for b in bodies_info if b["is_solid"]),
            "sheet_count": sum(1 for b in bodies_info if b["is_sheet"]),
            "bodies": bodies_info
        }

    except Exception as e:
        result["error"] = str(e)

    return result


def get_volume(part):
    """Get total volume of all solid bodies.

    NX Open API: MeasureManager.NewMassProperties()
    """
    result = {"success": False, "error": None, "data": {}}

    try:
        solid_bodies = get_solid_bodies(part)

        if not solid_bodies:
            result["error"] = "No solid bodies found in part"
            return result

        measure_manager = part.MeasureManager
        units = part.UnitCollection

        body_array = [NXOpen.IBody.Wrap(body) for body in solid_bodies]

        mass_props = measure_manager.NewMassProperties(
            units.FindObject("SquareMilliMeter"),
            units.FindObject("CubicMillimeter"),
            units.FindObject("Kilogram"),
            body_array
        )

        result["success"] = True
        result["data"] = {
            "volume": mass_props.Volume,
            "unit": "mm^3",
            "body_count": len(solid_bodies)
        }

    except Exception as e:
        result["error"] = str(e)

    return result


def get_surface_area(part):
    """Get total surface area of all solid bodies.

    NX Open API: MeasureManager.NewMassProperties()
    """
    result = {"success": False, "error": None, "data": {}}

    try:
        solid_bodies = get_solid_bodies(part)

        if not solid_bodies:
            result["error"] = "No solid bodies found in part"
            return result

        measure_manager = part.MeasureManager
        units = part.UnitCollection

        body_array = [NXOpen.IBody.Wrap(body) for body in solid_bodies]

        mass_props = measure_manager.NewMassProperties(
            units.FindObject("SquareMilliMeter"),
            units.FindObject("CubicMillimeter"),
            units.FindObject("Kilogram"),
            body_array
        )

        result["success"] = True
        result["data"] = {
            "surface_area": mass_props.Area,
            "unit": "mm^2",
            "body_count": len(solid_bodies)
        }

    except Exception as e:
        result["error"] = str(e)

    return result


def get_material(part):
    """Get material information from bodies in the part.

    NX Open API: Body.GetPhysicalMaterial()
    """
    result = {"success": False, "error": None, "data": {}}

    try:
        solid_bodies = get_solid_bodies(part)

        if not solid_bodies:
            result["error"] = "No solid bodies found in part"
            return result

        materials = {}

        for body in solid_bodies:
            try:
                phys_mat = body.GetPhysicalMaterial()
                if phys_mat:
                    mat_name = phys_mat.Name
                    if mat_name not in materials:
                        mat_data = {"name": mat_name}
                        # Try to get properties
                        try:
                            mat_data["density"] = phys_mat.GetRealPropertyValue("Density")
                        except:
                            pass
                        try:
                            mat_data["youngs_modulus"] = phys_mat.GetRealPropertyValue("YoungsModulus")
                        except:
                            pass
                        try:
                            mat_data["poissons_ratio"] = phys_mat.GetRealPropertyValue("PoissonsRatio")
                        except:
                            pass
                        materials[mat_name] = mat_data
            except:
                pass

        result["success"] = True
        result["data"] = {
            "material_count": len(materials),
            "materials": materials,
            "body_count": len(solid_bodies)
        }

    except Exception as e:
        result["error"] = str(e)

    return result


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


def _get_run_journal_exe() -> str:
    """Get the path to run_journal.exe."""
    return os.path.join(NX_BIN_PATH, "run_journal.exe")


def _run_journal(journal_path: str, args: list) -> Tuple[bool, str]:
    """Run an NX journal with arguments.

    Returns:
        Tuple of (success, output_or_error)
    """
    run_journal = _get_run_journal_exe()

    if not os.path.exists(run_journal):
        return False, f"run_journal.exe not found at {run_journal}"

    cmd = [run_journal, journal_path, "-args"] + args

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

        if result.returncode != 0:
            return False, f"Journal execution failed: {result.stderr}"

        return True, result.stdout

    except subprocess.TimeoutExpired:
        return False, "Journal execution timed out"
    except Exception as e:
        return False, str(e)


def _execute_geometry_operation(
    operation: str,
    part_path: str,
    extra_args: list = None
) -> Dict[str, Any]:
    """Execute a geometry query operation via NX journal.

    Args:
        operation: The operation to perform
        part_path: Path to the part file
        extra_args: Additional arguments for the operation

    Returns:
        Dict with operation result
    """
    # Create temporary journal file
    with tempfile.NamedTemporaryFile(
        mode='w',
        suffix='.py',
        delete=False
    ) as journal_file:
        journal_file.write(GEOMETRY_QUERY_JOURNAL)
        journal_path = journal_file.name

    # Create temporary output file
    output_file = tempfile.NamedTemporaryFile(
        mode='w',
        suffix='.json',
        delete=False
    ).name

    try:
        # Build arguments
        args = [operation, part_path, output_file]
        if extra_args:
            args.extend(extra_args)

        # Run the journal
        success, output = _run_journal(journal_path, args)

        if not success:
            return {"success": False, "error": output, "data": {}}

        # Read the result
        if os.path.exists(output_file):
            with open(output_file, 'r') as f:
                return json.load(f)
        else:
            return {"success": False, "error": "Output file not created", "data": {}}

    finally:
        # Cleanup temporary files
        if os.path.exists(journal_path):
            os.unlink(journal_path)
        if os.path.exists(output_file):
            os.unlink(output_file)


# =============================================================================
# Public API
# =============================================================================

def get_mass_properties(part_path: str) -> Dict[str, Any]:
    """Get mass properties from an NX part.

    Args:
        part_path: Full path to the .prt file

    Returns:
        Dict with keys:
            - success: bool
            - error: Optional error message
            - data: Dict with mass, volume, surface_area, centroid,
                    principal_moments, material, density, body_count

    Example:
        >>> result = get_mass_properties("C:/models/bracket.prt")
        >>> if result["success"]:
        ...     print(f"Mass: {result['data']['mass']} kg")
        ...     print(f"Volume: {result['data']['volume']} mm^3")
    """
    part_path = os.path.abspath(part_path)

    if not os.path.exists(part_path):
        return {
            "success": False,
            "error": f"Part file not found: {part_path}",
            "data": {}
        }

    return _execute_geometry_operation("mass_properties", part_path)


def get_bodies(part_path: str) -> Dict[str, Any]:
    """Get information about bodies in an NX part.

    Args:
        part_path: Full path to the .prt file

    Returns:
        Dict with keys:
            - success: bool
            - error: Optional error message
            - data: Dict with count, solid_count, sheet_count, bodies list

    Example:
        >>> result = get_bodies("C:/models/bracket.prt")
        >>> if result["success"]:
        ...     print(f"Solid bodies: {result['data']['solid_count']}")
    """
    part_path = os.path.abspath(part_path)

    if not os.path.exists(part_path):
        return {
            "success": False,
            "error": f"Part file not found: {part_path}",
            "data": {}
        }

    return _execute_geometry_operation("bodies", part_path)


def get_volume(part_path: str) -> Dict[str, Any]:
    """Get total volume of solid bodies in an NX part.

    Args:
        part_path: Full path to the .prt file

    Returns:
        Dict with keys:
            - success: bool
            - error: Optional error message
            - data: Dict with volume (mm^3), unit, body_count

    Example:
        >>> result = get_volume("C:/models/bracket.prt")
        >>> if result["success"]:
        ...     print(f"Volume: {result['data']['volume']} mm^3")
    """
    part_path = os.path.abspath(part_path)

    if not os.path.exists(part_path):
        return {
            "success": False,
            "error": f"Part file not found: {part_path}",
            "data": {}
        }

    return _execute_geometry_operation("volume", part_path)


def get_surface_area(part_path: str) -> Dict[str, Any]:
    """Get total surface area of solid bodies in an NX part.

    Args:
        part_path: Full path to the .prt file

    Returns:
        Dict with keys:
            - success: bool
            - error: Optional error message
            - data: Dict with surface_area (mm^2), unit, body_count

    Example:
        >>> result = get_surface_area("C:/models/bracket.prt")
        >>> if result["success"]:
        ...     print(f"Area: {result['data']['surface_area']} mm^2")
    """
    part_path = os.path.abspath(part_path)

    if not os.path.exists(part_path):
        return {
            "success": False,
            "error": f"Part file not found: {part_path}",
            "data": {}
        }

    return _execute_geometry_operation("surface_area", part_path)


def get_material(part_path: str) -> Dict[str, Any]:
    """Get material information from bodies in an NX part.

    Args:
        part_path: Full path to the .prt file

    Returns:
        Dict with keys:
            - success: bool
            - error: Optional error message
            - data: Dict with material_count, materials dict, body_count

    Example:
        >>> result = get_material("C:/models/bracket.prt")
        >>> if result["success"]:
        ...     for name, mat in result["data"]["materials"].items():
        ...         print(f"{name}: density={mat.get('density')}")
    """
    part_path = os.path.abspath(part_path)

    if not os.path.exists(part_path):
        return {
            "success": False,
            "error": f"Part file not found: {part_path}",
            "data": {}
        }

    return _execute_geometry_operation("material", part_path)