Atomizer/optimization_engine/extractors/extract_part_mass_material.py

"""
Extract Part Mass and Material Properties from NX Part Files

This extractor reads mass and material data from a temp file written by
the NX journal: nx_journals/extract_part_mass_material.py

The journal uses NXOpen.MeasureManager to extract:
- Mass (kg)
- Volume (mm^3)
- Surface area (mm^2)
- Center of gravity (mm)
- Material name and density

NX Open APIs Used (by journal):
    - NXOpen.MeasureManager.NewMassProperties()
    - NXOpen.MeasureBodies
    - NXOpen.Body.GetBodies()
    - NXOpen.PhysicalMaterial

Author: Atomizer
Created: 2025-12-05
Version: 1.0
"""

from pathlib import Path
from typing import Dict, Any, Optional, Union
import json


def extract_part_mass_material(
    prt_file: Union[str, Path],
    properties_file: Optional[Union[str, Path]] = None
) -> Dict[str, Any]:
    """
    Extract mass and material properties from NX part file.

    This function reads from a temp JSON file that must be created by
    running the NX journal: nx_journals/extract_part_mass_material.py

    Args:
        prt_file: Path to .prt file (used to locate temp file)
        properties_file: Optional explicit path to _temp_part_properties.json
                        If not provided, looks in same directory as prt_file

    Returns:
        Dictionary containing:
        - 'mass_kg': Mass in kilograms (float)
        - 'mass_g': Mass in grams (float)
        - 'volume_mm3': Volume in mm^3 (float)
        - 'surface_area_mm2': Surface area in mm^2 (float)
        - 'center_of_gravity_mm': [x, y, z] in mm (list)
        - 'moments_of_inertia': {'Ixx', 'Iyy', 'Izz', 'unit'} or None
        - 'material': {'name', 'density', 'density_unit'} (dict)
        - 'num_bodies': Number of solid bodies (int)

    Raises:
        FileNotFoundError: If prt file or temp properties file not found
        ValueError: If temp file has invalid format or extraction failed

    Example:
        >>> result = extract_part_mass_material('model.prt')
        >>> print(f"Mass: {result['mass_kg']:.3f} kg")
        Mass: 1.234 kg
        >>> print(f"Material: {result['material']['name']}")
        Material: Aluminum_6061

    Note:
        Before calling this function, you must run the NX journal to
        create the temp file:
        ```
        run_journal.exe extract_part_mass_material.py model.prt
        ```
    """
    prt_file = Path(prt_file)

    if not prt_file.exists():
        raise FileNotFoundError(f"Part file not found: {prt_file}")

    # Determine properties file location
    if properties_file:
        props_file = Path(properties_file)
    else:
        props_file = prt_file.parent / "_temp_part_properties.json"

    if not props_file.exists():
        raise FileNotFoundError(
            f"Part properties temp file not found: {props_file}\n"
            f"Run the NX journal first:\n"
            f"  run_journal.exe extract_part_mass_material.py {prt_file}"
        )

    # Read and parse JSON
    try:
        with open(props_file, 'r') as f:
            data = json.load(f)
    except json.JSONDecodeError as e:
        raise ValueError(f"Invalid JSON in properties file: {e}")

    # Check for extraction errors
    if not data.get('success', False):
        error_msg = data.get('error', 'Unknown error during extraction')
        raise ValueError(f"NX extraction failed: {error_msg}")

    # Build result dictionary
    result = {
        'mass_kg': float(data.get('mass_kg', 0.0)),
        'mass_g': float(data.get('mass_g', 0.0)),
        'volume_mm3': float(data.get('volume_mm3', 0.0)),
        'surface_area_mm2': float(data.get('surface_area_mm2', 0.0)),
        'center_of_gravity_mm': data.get('center_of_gravity_mm', [0.0, 0.0, 0.0]),
        'moments_of_inertia': data.get('moments_of_inertia'),
        'material': data.get('material', {'name': None, 'density': None, 'density_unit': 'kg/mm^3'}),
        'num_bodies': int(data.get('num_bodies', 0)),
        'part_file': data.get('part_file', prt_file.name),
    }

    print(f"[OK] Part mass: {result['mass_kg']:.6f} kg ({result['mass_g']:.2f} g)")
    if result['material'].get('name'):
        print(f"[OK] Material: {result['material']['name']}")

    return result


def extract_part_mass(
    prt_file: Union[str, Path],
    properties_file: Optional[Union[str, Path]] = None
) -> float:
    """
    Convenience function to extract just the mass in kg.

    Args:
        prt_file: Path to .prt file
        properties_file: Optional explicit path to temp file

    Returns:
        Mass in kilograms (float)

    Example:
        >>> mass = extract_part_mass('model.prt')
        >>> print(f"Mass: {mass:.3f} kg")
        Mass: 1.234 kg
    """
    result = extract_part_mass_material(prt_file, properties_file)
    return result['mass_kg']


def extract_part_material(
    prt_file: Union[str, Path],
    properties_file: Optional[Union[str, Path]] = None
) -> Dict[str, Any]:
    """
    Convenience function to extract just material info.

    Args:
        prt_file: Path to .prt file
        properties_file: Optional explicit path to temp file

    Returns:
        Dictionary with 'name', 'density', 'density_unit'

    Example:
        >>> mat = extract_part_material('model.prt')
        >>> print(f"Material: {mat['name']}, Density: {mat['density']}")
        Material: Steel_304, Density: 7.93e-06
    """
    result = extract_part_mass_material(prt_file, properties_file)
    return result['material']


class PartMassExtractor:
    """
    Class-based extractor for part mass and material with caching.

    Use this when you need to extract properties from multiple parts
    or want to cache results.

    Example:
        >>> extractor = PartMassExtractor('model.prt')
        >>> result = extractor.extract()
        >>> print(result['mass_kg'])
        1.234
    """

    def __init__(
        self,
        prt_file: Union[str, Path],
        properties_file: Optional[Union[str, Path]] = None
    ):
        """
        Initialize the extractor.

        Args:
            prt_file: Path to .prt file
            properties_file: Optional explicit path to temp file
        """
        self.prt_file = Path(prt_file)
        self.properties_file = Path(properties_file) if properties_file else None
        self._cached_result = None

    def extract(self, use_cache: bool = True) -> Dict[str, Any]:
        """
        Extract mass and material properties.

        Args:
            use_cache: If True, returns cached result if available

        Returns:
            Dictionary with all extracted properties
        """
        if use_cache and self._cached_result is not None:
            return self._cached_result

        self._cached_result = extract_part_mass_material(
            self.prt_file,
            self.properties_file
        )
        return self._cached_result

    @property
    def mass_kg(self) -> float:
        """Get mass in kg (extracts if needed)."""
        return self.extract()['mass_kg']

    @property
    def mass_g(self) -> float:
        """Get mass in grams (extracts if needed)."""
        return self.extract()['mass_g']

    @property
    def material_name(self) -> Optional[str]:
        """Get material name (extracts if needed)."""
        return self.extract()['material'].get('name')

    @property
    def density(self) -> Optional[float]:
        """Get material density in kg/mm^3 (extracts if needed)."""
        return self.extract()['material'].get('density')

    def clear_cache(self):
        """Clear the cached result."""
        self._cached_result = None


if __name__ == "__main__":
    import sys

    if len(sys.argv) < 2:
        print(f"Usage: python {sys.argv[0]} <prt_file> [properties_file]")
        sys.exit(1)

    prt_file = Path(sys.argv[1])
    props_file = Path(sys.argv[2]) if len(sys.argv) > 2 else None

    try:
        result = extract_part_mass_material(prt_file, props_file)
        print("\n" + "="*50)
        print("PART MASS & MATERIAL EXTRACTION RESULTS")
        print("="*50)
        print(f"Part File: {result['part_file']}")
        print(f"Mass: {result['mass_kg']:.6f} kg ({result['mass_g']:.2f} g)")
        print(f"Volume: {result['volume_mm3']:.2f} mm^3")
        print(f"Surface Area: {result['surface_area_mm2']:.2f} mm^2")
        print(f"Center of Gravity: {result['center_of_gravity_mm']} mm")
        print(f"Num Bodies: {result['num_bodies']}")
        if result['material']['name']:
            print(f"Material: {result['material']['name']}")
            if result['material']['density']:
                print(f"Density: {result['material']['density']} {result['material']['density_unit']}")
        if result['moments_of_inertia']:
            moi = result['moments_of_inertia']
            print(f"Moments of Inertia: Ixx={moi['Ixx']}, Iyy={moi['Iyy']}, Izz={moi['Izz']} {moi['unit']}")
    except Exception as e:
        print(f"\nERROR: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)