Atomizer/optimization_engine/extractors/introspect_part.py

"""
Part Introspection Extractor (E12)
===================================

Comprehensive introspection of NX .prt files. Extracts:
- All expressions (user and internal)
- Mass properties (mass, volume, surface area, CoG)
- Material properties
- Body information
- Part attributes
- Groups
- Features summary
- Datum planes and coordinate systems
- Unit system
- Linked/associated parts

Usage:
    from optimization_engine.extractors.introspect_part import introspect_part

    result = introspect_part("path/to/model.prt")
    print(f"Mass: {result['mass_properties']['mass_kg']} kg")
    print(f"Expressions: {result['expressions']['user_count']}")

Dependencies:
    - NX installed with run_journal.exe
    - SPLM_LICENSE_SERVER environment variable

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

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


# NX installation path
NX_INSTALL_PATH = r"C:\Program Files\Siemens\DesigncenterNX2512"
RUN_JOURNAL_EXE = os.path.join(NX_INSTALL_PATH, "NXBIN", "run_journal.exe")
JOURNAL_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))),
                             "nx_journals", "introspect_part.py")

# License server
LICENSE_SERVER = "28000@dalidou;28000@100.80.199.40"


def introspect_part(
    prt_file_path: str,
    output_dir: Optional[str] = None,
    timeout_seconds: int = 120,
    verbose: bool = True
) -> Dict[str, Any]:
    """
    Run comprehensive introspection on an NX .prt file.

    Args:
        prt_file_path: Path to the .prt file
        output_dir: Directory for output JSON (defaults to prt directory)
        timeout_seconds: Maximum time to wait for journal execution
        verbose: Print progress messages

    Returns:
        Dictionary with introspection results:
        {
            'success': bool,
            'part_file': str,
            'expressions': {
                'user': [...],
                'internal': [...],
                'user_count': int,
                'total_count': int
            },
            'mass_properties': {
                'mass_kg': float,
                'mass_g': float,
                'volume_mm3': float,
                'surface_area_mm2': float,
                'center_of_gravity_mm': [x, y, z]
            },
            'materials': {
                'assigned': [...],
                'available': [...]
            },
            'bodies': {
                'solid_bodies': [...],
                'sheet_bodies': [...],
                'counts': {...}
            },
            'attributes': [...],
            'groups': [...],
            'features': {
                'total_count': int,
                'by_type': {...}
            },
            'datums': {...},
            'units': {...},
            'linked_parts': {...}
        }

    Raises:
        FileNotFoundError: If prt file doesn't exist
        RuntimeError: If journal execution fails
    """
    prt_path = Path(prt_file_path)
    if not prt_path.exists():
        raise FileNotFoundError(f"Part file not found: {prt_file_path}")

    if output_dir is None:
        output_dir = str(prt_path.parent)

    output_file = os.path.join(output_dir, "_temp_introspection.json")

    # Remove old output file if exists
    if os.path.exists(output_file):
        os.remove(output_file)

    if verbose:
        print(f"[INTROSPECT] Part: {prt_path.name}")
        print(f"[INTROSPECT] Output: {output_dir}")

    # Build PowerShell command (learned workaround - see LAC)
    # Using [Environment]::SetEnvironmentVariable() for reliable license server setting
    ps_command = (
        f"[Environment]::SetEnvironmentVariable('SPLM_LICENSE_SERVER', '{LICENSE_SERVER}', 'Process'); "
        f"& '{RUN_JOURNAL_EXE}' '{JOURNAL_PATH}' -args '{prt_file_path}' '{output_dir}' 2>&1"
    )

    cmd = ["powershell", "-Command", ps_command]

    if verbose:
        print(f"[INTROSPECT] Executing NX journal...")

    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=timeout_seconds
        )

        if verbose and result.stdout:
            # Print key lines from output
            for line in result.stdout.split('\n'):
                if '[INTROSPECT]' in line:
                    print(line)

        if result.returncode != 0 and verbose:
            print(f"[INTROSPECT] Warning: Non-zero return code: {result.returncode}")
            if result.stderr:
                print(f"[INTROSPECT] Stderr: {result.stderr[:500]}")

    except subprocess.TimeoutExpired:
        raise RuntimeError(f"Journal execution timed out after {timeout_seconds}s")
    except Exception as e:
        raise RuntimeError(f"Journal execution failed: {e}")

    # Wait for output file (NX may still be writing)
    wait_start = time.time()
    while not os.path.exists(output_file) and (time.time() - wait_start) < 10:
        time.sleep(0.5)

    if not os.path.exists(output_file):
        raise RuntimeError(f"Introspection output file not created: {output_file}")

    # Read and return results
    with open(output_file, 'r') as f:
        results = json.load(f)

    if verbose:
        if results.get('success'):
            print(f"[INTROSPECT] Success!")
            print(f"[INTROSPECT]   Expressions: {results['expressions']['user_count']} user")
            print(f"[INTROSPECT]   Mass: {results['mass_properties']['mass_kg']:.4f} kg")
        else:
            print(f"[INTROSPECT] Failed: {results.get('error', 'Unknown error')}")

    return results


def get_expressions_dict(introspection_result: Dict[str, Any]) -> Dict[str, float]:
    """
    Convert introspection result to simple {name: value} dictionary of user expressions.

    Args:
        introspection_result: Result from introspect_part()

    Returns:
        Dictionary mapping expression names to values
    """
    expressions = {}
    for expr in introspection_result.get('expressions', {}).get('user', []):
        name = expr.get('name')
        value = expr.get('value')
        if name and value is not None:
            expressions[name] = value
    return expressions


def get_expression_value(introspection_result: Dict[str, Any], name: str) -> Optional[float]:
    """
    Get a specific expression value from introspection result.

    Args:
        introspection_result: Result from introspect_part()
        name: Expression name

    Returns:
        Expression value or None if not found
    """
    for expr in introspection_result.get('expressions', {}).get('user', []):
        if expr.get('name') == name:
            return expr.get('value')
    return None


def print_introspection_summary(result: Dict[str, Any]) -> None:
    """Print a formatted summary of introspection results."""
    print("\n" + "="*60)
    print(f"INTROSPECTION SUMMARY: {result.get('part_file', 'Unknown')}")
    print("="*60)

    # Mass
    mp = result.get('mass_properties', {})
    print(f"\nMASS PROPERTIES:")
    print(f"  Mass: {mp.get('mass_kg', 0):.4f} kg ({mp.get('mass_g', 0):.2f} g)")
    print(f"  Volume: {mp.get('volume_mm3', 0):.2f} mm³")
    print(f"  Surface Area: {mp.get('surface_area_mm2', 0):.2f} mm²")
    cog = mp.get('center_of_gravity_mm', [0, 0, 0])
    print(f"  Center of Gravity: [{cog[0]:.2f}, {cog[1]:.2f}, {cog[2]:.2f}] mm")

    # Bodies
    bodies = result.get('bodies', {})
    counts = bodies.get('counts', {})
    print(f"\nBODIES:")
    print(f"  Solid: {counts.get('solid', 0)}")
    print(f"  Sheet: {counts.get('sheet', 0)}")

    # Materials
    mats = result.get('materials', {})
    print(f"\nMATERIALS:")
    for mat in mats.get('assigned', []):
        print(f"  {mat.get('name', 'Unknown')}")
        props = mat.get('properties', {})
        if 'Density' in props:
            print(f"    Density: {props['Density']}")
        if 'YoungModulus' in props:
            print(f"    Young's Modulus: {props['YoungModulus']}")

    # Expressions
    exprs = result.get('expressions', {})
    print(f"\nEXPRESSIONS: {exprs.get('user_count', 0)} user, {len(exprs.get('internal', []))} internal")
    user_exprs = exprs.get('user', [])
    if user_exprs:
        print("  User expressions (first 10):")
        for expr in user_exprs[:10]:
            units = f" [{expr.get('units', '')}]" if expr.get('units') else ""
            print(f"    {expr.get('name')}: {expr.get('value')}{units}")
        if len(user_exprs) > 10:
            print(f"    ... and {len(user_exprs) - 10} more")

    # Features
    feats = result.get('features', {})
    print(f"\nFEATURES: {feats.get('total_count', 0)} total")
    by_type = feats.get('by_type', {})
    if by_type:
        top_types = sorted(by_type.items(), key=lambda x: x[1], reverse=True)[:5]
        for feat_type, count in top_types:
            print(f"  {feat_type}: {count}")

    # Units
    units = result.get('units', {})
    print(f"\nUNITS: {units.get('system', 'Unknown')}")

    print("="*60 + "\n")


# CLI interface
if __name__ == "__main__":
    import sys

    if len(sys.argv) < 2:
        print("Usage: python introspect_part.py <prt_file_path> [output_dir]")
        sys.exit(1)

    prt_path = sys.argv[1]
    out_dir = sys.argv[2] if len(sys.argv) > 2 else None

    try:
        result = introspect_part(prt_path, out_dir, verbose=True)
        print_introspection_summary(result)
    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)