"""
NX Journal Script to Extract Part Mass and Material Properties

This script extracts:
- Mass (kg)
- Volume (mm^3)
- Surface area (mm^2)
- Center of gravity (mm)
- Material name (if assigned)
- Density (kg/mm^3)

Results are written to _temp_part_properties.json in the working directory.

Usage:
    run_journal.exe extract_part_mass_material.py <prt_file_path> [output_dir]

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

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

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


def get_all_solid_bodies(part):
    """Get all solid bodies from the part."""
    bodies = []
    try:
        # Get bodies from the part
        for body in part.Bodies:
            if body.IsSolidBody:
                bodies.append(body)
    except Exception as e:
        print(f"[JOURNAL] Warning getting bodies: {e}")
    return bodies


def get_material_info(part, body):
    """Extract material information from a body if assigned."""
    material_info = {
        'name': None,
        'density': None,
        'density_unit': 'kg/mm^3'
    }

    # Method 1: Try to get physical material directly from body
    try:
        phys_mat = body.GetPhysicalMaterial()
        if phys_mat:
            material_info['name'] = phys_mat.Name
            print(f"[JOURNAL]   Material from body: {phys_mat.Name}")

            # Try to get density property
            try:
                density_prop = phys_mat.GetPropertyValue("Density")
                if density_prop:
                    material_info['density'] = float(density_prop)
                    print(f"[JOURNAL]   Density: {density_prop}")
            except Exception as de:
                print(f"[JOURNAL]   Could not get density: {de}")
    except Exception as e:
        print(f"[JOURNAL]   GetPhysicalMaterial failed: {e}")

    # Method 2: If no material from body, try part-level material assignment
    if material_info['name'] is None:
        try:
            # Check if part has PhysicalMaterialManager
            pmm = part.PhysicalMaterialManager
            if pmm:
                # Get all materials in the part
                materials = pmm.GetAllPhysicalMaterials()
                if materials and len(materials) > 0:
                    # Use the first material as default
                    mat = materials[0]
                    material_info['name'] = mat.Name
                    print(f"[JOURNAL]   Material from PhysicalMaterialManager: {mat.Name}")

                    try:
                        density_prop = mat.GetPropertyValue("Density")
                        if density_prop:
                            material_info['density'] = float(density_prop)
                    except:
                        pass
        except Exception as e:
            print(f"[JOURNAL]   PhysicalMaterialManager failed: {e}")

    # Method 3: Try using body attributes
    if material_info['name'] is None:
        try:
            # Some NX versions store material as an attribute
            attrs = body.GetUserAttributes()
            for attr in attrs:
                if 'material' in attr.Title.lower():
                    material_info['name'] = attr.StringValue
                    print(f"[JOURNAL]   Material from attribute: {attr.StringValue}")
                    break
        except Exception as e:
            print(f"[JOURNAL]   Body attributes failed: {e}")

    return material_info


def extract_mass_properties(theSession, part, bodies):
    """
    Extract mass properties using MeasureManager.

    API Reference: https://www.nxjournaling.com/content/mass-properties-using-python

    Returns dict with mass, volume, area, center of gravity.
    """
    results = {
        'mass_kg': 0.0,
        'volume_mm3': 0.0,
        'surface_area_mm2': 0.0,
        'center_of_gravity_mm': [0.0, 0.0, 0.0],
        'moments_of_inertia': None,
        'num_bodies': len(bodies)
    }

    if not bodies:
        print("[JOURNAL] No solid bodies found in part")
        return results

    try:
        # Get the measure manager
        measureManager = part.MeasureManager

        # Convert bodies list to array for NX API
        bodyArray = bodies if isinstance(bodies, list) else list(bodies)

        # Get unit collection and build mass_units array
        # API requires: [Area, Volume, Mass, Length] base units
        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)
        measureBodies = measureManager.NewMassProperties(mass_units, 0.99, bodyArray)
        print("[JOURNAL] Using NewMassProperties(mass_units, accuracy, bodies) API")

        # Get the results
        if measureBodies:
            # Mass
            try:
                results['mass_kg'] = measureBodies.Mass
                print(f"[JOURNAL]   Raw mass value: {measureBodies.Mass}")
            except AttributeError as e:
                print(f"[JOURNAL]   Mass attribute error: {e}")

            # Volume
            try:
                results['volume_mm3'] = measureBodies.Volume
            except AttributeError:
                pass

            # Surface area
            try:
                results['surface_area_mm2'] = measureBodies.Area
            except AttributeError:
                pass

            # Center of gravity
            try:
                cog = measureBodies.Centroid
                if cog:
                    results['center_of_gravity_mm'] = [cog.X, cog.Y, cog.Z]
            except AttributeError:
                pass

            # Moments of inertia - NewMassProperties doesn't provide inertia tensors
            # Would need different API for that
            results['moments_of_inertia'] = None

            # Dispose
            try:
                measureBodies.Dispose()
            except:
                pass

    except Exception as e:
        print(f"[JOURNAL] Error in mass properties extraction: {e}")
        import traceback
        traceback.print_exc()

    return results


def main(args):
    """
    Main entry point for NX journal.

    Args:
        args[0]: .prt file path
        args[1]: output directory (optional, defaults to prt directory)
    """
    if len(args) < 1:
        print("ERROR: No .prt file path provided")
        print("Usage: run_journal.exe extract_part_mass_material.py <prt_file> [output_dir]")
        return False

    prt_file_path = args[0]
    output_dir = args[1] if len(args) > 1 else os.path.dirname(prt_file_path)

    prt_filename = os.path.basename(prt_file_path)

    print(f"[JOURNAL] " + "="*60)
    print(f"[JOURNAL] NX PART MASS & MATERIAL EXTRACTOR")
    print(f"[JOURNAL] " + "="*60)
    print(f"[JOURNAL] Part file: {prt_filename}")
    print(f"[JOURNAL] Output dir: {output_dir}")

    results = {
        'part_file': prt_filename,
        'part_path': prt_file_path,
        'mass_kg': 0.0,
        'mass_g': 0.0,
        'volume_mm3': 0.0,
        'surface_area_mm2': 0.0,
        'center_of_gravity_mm': [0.0, 0.0, 0.0],
        'moments_of_inertia': None,
        'material': {
            'name': None,
            'density': None,
            'density_unit': 'kg/mm^3'
        },
        'num_bodies': 0,
        'success': False,
        'error': None
    }

    try:
        theSession = NXOpen.Session.GetSession()

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

        # Open the part file
        print(f"[JOURNAL] Opening part file...")
        basePart, partLoadStatus = theSession.Parts.OpenActiveDisplay(
            prt_file_path,
            NXOpen.DisplayPartOption.AllowAdditional
        )
        partLoadStatus.Dispose()

        workPart = theSession.Parts.Work
        print(f"[JOURNAL] Loaded part: {workPart.Name}")

        # Get all solid bodies
        bodies = get_all_solid_bodies(workPart)
        print(f"[JOURNAL] Found {len(bodies)} solid bodies")

        if bodies:
            # Extract mass properties
            mass_props = extract_mass_properties(theSession, workPart, bodies)
            results.update(mass_props)
            results['mass_g'] = results['mass_kg'] * 1000.0

            # Get material from first body (typically all bodies have same material)
            material_info = get_material_info(workPart, bodies[0])
            results['material'] = material_info

            results['success'] = True

            print(f"[JOURNAL] ")
            print(f"[JOURNAL] RESULTS:")
            print(f"[JOURNAL]   Mass: {results['mass_kg']:.6f} kg ({results['mass_g']:.2f} g)")
            print(f"[JOURNAL]   Volume: {results['volume_mm3']:.2f} mm^3")
            print(f"[JOURNAL]   Surface Area: {results['surface_area_mm2']:.2f} mm^2")
            print(f"[JOURNAL]   CoG: [{results['center_of_gravity_mm'][0]:.2f}, {results['center_of_gravity_mm'][1]:.2f}, {results['center_of_gravity_mm'][2]:.2f}] mm")
            if material_info['name']:
                print(f"[JOURNAL]   Material: {material_info['name']}")
                if material_info['density']:
                    print(f"[JOURNAL]   Density: {material_info['density']} {material_info['density_unit']}")
        else:
            results['error'] = "No solid bodies found in part"
            print(f"[JOURNAL] ERROR: No solid bodies found")

        # Write results to JSON file
        output_file = os.path.join(output_dir, "_temp_part_properties.json")
        with open(output_file, 'w') as f:
            json.dump(results, f, indent=2)
        print(f"[JOURNAL] Results written to: {output_file}")

        # Also write simple mass value for backward compatibility
        mass_file = os.path.join(output_dir, "_temp_mass.txt")
        with open(mass_file, 'w') as f:
            f.write(str(results['mass_kg']))
        print(f"[JOURNAL] Mass written to: {mass_file}")

        return True

    except Exception as e:
        results['error'] = str(e)
        results['success'] = False
        print(f"[JOURNAL] FATAL ERROR: {e}")
        import traceback
        traceback.print_exc()

        # Still write results file even on error
        output_file = os.path.join(output_dir, "_temp_part_properties.json")
        try:
            with open(output_file, 'w') as f:
                json.dump(results, f, indent=2)
        except:
            pass

        return False


if __name__ == '__main__':
    main(sys.argv[1:])