Atomizer/optimization_engine/nx/solve_simulation.py

"""
NX Journal Script to Solve Simulation in Batch Mode

This script handles BOTH single-part simulations AND multi-part assembly FEMs.

=============================================================================
MULTI-PART ASSEMBLY FEM WORKFLOW (for .afm-based simulations)
=============================================================================

Based on recorded NX journal from interactive session (Nov 28, 2025).

The correct workflow for assembly FEM updates:

1. LOAD PARTS
   - Open ASSY_M1.prt and M1_Blank_fem1_i.prt to have geometry loaded
   - Find and switch to M1_Blank part for expression editing

2. UPDATE EXPRESSIONS
   - Switch to modeling application
   - Edit expressions with units
   - Call MakeUpToDate() on modified expressions
   - Call DoUpdate() to rebuild geometry

3. SWITCH TO SIM AND UPDATE FEM COMPONENTS
   - Open the .sim file
   - Navigate component hierarchy via RootComponent.FindObject()
   - For each component FEM:
     - SetWorkComponent() to make it the work part
     - FindObject("FEModel").UpdateFemodel()

4. MERGE DUPLICATE NODES (critical for assembly FEM!)
   - Switch to assembly FEM component
   - CreateDuplicateNodesCheckBuilder()
   - Set MergeOccurrenceNodes = True
   - IdentifyDuplicateNodes() then MergeDuplicateNodes()

5. RESOLVE LABEL CONFLICTS
   - CreateAssemblyLabelManagerBuilder()
   - SetFEModelOccOffsets() for each occurrence
   - Commit()

6. SOLVE
   - SetWorkComponent(Null) to return to sim level
   - SolveChainOfSolutions()

=============================================================================
"""

import sys
import os
import NXOpen
import NXOpen.Assemblies
import NXOpen.CAE


def extract_part_mass(theSession, part, output_dir):
    """
    Extract mass from a part using NX MeasureManager.

    Writes mass to _temp_mass.txt and _temp_part_properties.json in output_dir.

    Args:
        theSession: NXOpen.Session
        part: NXOpen.Part to extract mass from
        output_dir: Directory to write temp files

    Returns:
        Mass in kg (float)
    """
    import json

    results = {
        "part_file": part.Name,
        "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],
        "num_bodies": 0,
        "success": False,
        "error": None,
    }

    try:
        # Get all solid bodies
        bodies = []
        for body in part.Bodies:
            if body.IsSolidBody:
                bodies.append(body)

        results["num_bodies"] = len(bodies)

        if not bodies:
            results["error"] = "No solid bodies found"
            raise ValueError("No solid bodies found in part")

        # Get the measure manager
        measureManager = part.MeasureManager

        # 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
        measureBodies = measureManager.NewMassProperties(mass_units, 0.99, bodies)

        if measureBodies:
            results["mass_kg"] = measureBodies.Mass
            results["mass_g"] = results["mass_kg"] * 1000.0

            try:
                results["volume_mm3"] = measureBodies.Volume
            except:
                pass

            try:
                results["surface_area_mm2"] = measureBodies.Area
            except:
                pass

            try:
                cog = measureBodies.Centroid
                if cog:
                    results["center_of_gravity_mm"] = [cog.X, cog.Y, cog.Z]
            except:
                pass

            try:
                measureBodies.Dispose()
            except:
                pass

        results["success"] = True

    except Exception as e:
        results["error"] = str(e)
        results["success"] = False

    # 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)

    # 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"]))

    if not results["success"]:
        raise ValueError(results["error"])

    return results["mass_kg"]


def find_or_open_part(theSession, part_path):
    """
    Find a part if already loaded, otherwise open it.
    In NX, calling Parts.Open() on an already-loaded part raises 'File already exists'.
    """
    part_name = os.path.splitext(os.path.basename(part_path))[0]

    # Try to find in already-loaded parts
    for part in theSession.Parts:
        if part.Name == part_name:
            return part, True
        try:
            if part.FullPath and part.FullPath.lower() == part_path.lower():
                return part, True
        except:
            pass

    # Not found, open it
    markId = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Visible, f"Load {part_name}")
    part, partLoadStatus = theSession.Parts.Open(part_path)
    partLoadStatus.Dispose()
    return part, False


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

    Args:
        args: Command line arguments
              args[0]: .sim file path
              args[1]: solution_name (optional, or "None" for default)
              args[2+]: expression updates as "name=value" pairs
    """
    if len(args) < 1:
        print("ERROR: No .sim file path provided")
        print(
            "Usage: run_journal.exe solve_simulation.py <sim_file_path> [solution_name] [expr1=val1] ..."
        )
        return False

    sim_file_path = args[0]
    solution_name = args[1] if len(args) > 1 and args[1] != "None" else None

    # Parse expression updates
    expression_updates = {}
    for arg in args[2:]:
        if "=" in arg:
            name, value = arg.split("=", 1)
            expression_updates[name] = float(value)

    # Resolve ALL paths to absolute — NX can fail silently with ".." in paths
    sim_file_path = os.path.abspath(sim_file_path)
    working_dir = os.path.dirname(sim_file_path)
    sim_filename = os.path.basename(sim_file_path)

    print(f"[JOURNAL] " + "=" * 60)
    print(f"[JOURNAL] NX SIMULATION SOLVER (Assembly FEM Workflow)")
    print(f"[JOURNAL] " + "=" * 60)
    print(f"[JOURNAL] Simulation: {sim_filename}")
    print(f"[JOURNAL] Working directory: {working_dir}")
    print(f"[JOURNAL] Solution: {solution_name or 'Solution 1'}")
    print(f"[JOURNAL] Expression updates: {len(expression_updates)}")
    for name, value in expression_updates.items():
        print(f"[JOURNAL]   {name} = {value}")

    try:
        theSession = NXOpen.Session.GetSession()

        # Set load options
        theSession.Parts.LoadOptions.LoadLatest = False
        theSession.Parts.LoadOptions.ComponentLoadMethod = (
            NXOpen.LoadOptions.LoadMethod.FromDirectory
        )
        theSession.Parts.LoadOptions.SetSearchDirectories([working_dir], [True])
        theSession.Parts.LoadOptions.ComponentsToLoad = NXOpen.LoadOptions.LoadComponents.All
        theSession.Parts.LoadOptions.PartLoadOption = NXOpen.LoadOptions.LoadOption.FullyLoad
        theSession.Parts.LoadOptions.SetInterpartData(True, NXOpen.LoadOptions.Parent.All)
        theSession.Parts.LoadOptions.AbortOnFailure = False

        # Close any open parts
        try:
            theSession.Parts.CloseAll([NXOpen.BasePart.CloseWholeTree])
        except:
            pass

        # Check for assembly FEM files
        afm_files = [f for f in os.listdir(working_dir) if f.endswith(".afm")]
        is_assembly = len(afm_files) > 0

        if is_assembly and expression_updates:
            print(f"[JOURNAL] ")
            print(f"[JOURNAL] DETECTED: Multi-part Assembly FEM")
            print(f"[JOURNAL] Using ASSEMBLY FEM WORKFLOW")
            print(f"[JOURNAL] ")
            return solve_assembly_fem_workflow(
                theSession, sim_file_path, solution_name, expression_updates, working_dir
            )
        else:
            print(f"[JOURNAL] ")
            print(f"[JOURNAL] Using SIMPLE WORKFLOW (no expression updates or single-part)")
            print(f"[JOURNAL] ")
            return solve_simple_workflow(
                theSession, sim_file_path, solution_name, expression_updates, working_dir
            )

    except Exception as e:
        print(f"[JOURNAL] FATAL ERROR: {e}")
        import traceback

        traceback.print_exc()
        return False


def solve_assembly_fem_workflow(
    theSession, sim_file_path, solution_name, expression_updates, working_dir
):
    """
    Full assembly FEM workflow based on recorded NX journal.

    This is the correct workflow for multi-part assembly FEMs.
    """
    sim_filename = os.path.basename(sim_file_path)

    # ==========================================================================
    # STEP 1: LOAD SIM FILE FIRST (loads entire assembly hierarchy)
    # ==========================================================================
    print(f"[JOURNAL] STEP 1: Loading SIM file and all components...")

    # CRITICAL: Open the SIM file FIRST using OpenActiveDisplay
    # This loads the entire assembly FEM hierarchy (.afm, .fem, associated .prt files)
    # The sim file knows its component structure and will load everything it needs
    sim_file_full_path = os.path.join(working_dir, sim_filename)
    print(f"[JOURNAL]   Opening SIM file: {sim_filename}")
    basePart, partLoadStatus = theSession.Parts.OpenActiveDisplay(
        sim_file_full_path, NXOpen.DisplayPartOption.AllowAdditional
    )
    partLoadStatus.Dispose()

    workSimPart = theSession.Parts.BaseWork
    displaySimPart = theSession.Parts.BaseDisplay
    print(f"[JOURNAL]   SIM loaded: {workSimPart.Name}")

    # List loaded parts
    print(f"[JOURNAL]   Currently loaded parts:")
    for part in theSession.Parts:
        print(f"[JOURNAL]     - {part.Name}")

    # ==========================================================================
    # STEP 1b: LOAD GEOMETRY PARTS FOR EXPRESSION EDITING
    # ==========================================================================
    print(f"[JOURNAL] STEP 1b: Loading geometry parts for expression editing...")

    # The recorded journal loads these geometry parts explicitly:
    # 1. ASSY_M1.prt - the main geometry assembly
    # 2. M1_Blank_fem1_i.prt - idealized geometry for M1_Blank FEM
    # 3. M1_Vertical_Support_Skeleton_fem1_i.prt - idealized geometry for support FEM

    # Load ASSY_M1.prt (to have the geometry assembly available)
    assy_prt_path = os.path.join(working_dir, "ASSY_M1.prt")
    if os.path.exists(assy_prt_path):
        print(f"[JOURNAL]   Loading ASSY_M1.prt...")
        part1, was_loaded = find_or_open_part(theSession, assy_prt_path)
        if was_loaded:
            print(f"[JOURNAL]     (already loaded)")
    else:
        print(f"[JOURNAL]   WARNING: ASSY_M1.prt not found!")

    # Load M1_Blank_fem1_i.prt (idealized geometry for M1_Blank)
    idealized_prt_path = os.path.join(working_dir, "M1_Blank_fem1_i.prt")
    if os.path.exists(idealized_prt_path):
        print(f"[JOURNAL]   Loading M1_Blank_fem1_i.prt...")
        part2, was_loaded = find_or_open_part(theSession, idealized_prt_path)
        if was_loaded:
            print(f"[JOURNAL]     (already loaded)")
    else:
        print(f"[JOURNAL]   WARNING: M1_Blank_fem1_i.prt not found!")

    # Load M1_Vertical_Support_Skeleton_fem1_i.prt (CRITICAL: idealized geometry for support)
    skeleton_idealized_prt_path = os.path.join(
        working_dir, "M1_Vertical_Support_Skeleton_fem1_i.prt"
    )
    if os.path.exists(skeleton_idealized_prt_path):
        print(f"[JOURNAL]   Loading M1_Vertical_Support_Skeleton_fem1_i.prt...")
        part3_skel, was_loaded = find_or_open_part(theSession, skeleton_idealized_prt_path)
        if was_loaded:
            print(f"[JOURNAL]     (already loaded)")
    else:
        print(f"[JOURNAL]   WARNING: M1_Vertical_Support_Skeleton_fem1_i.prt not found!")

    # ==========================================================================
    # STEP 2: UPDATE EXPRESSIONS IN M1_BLANK AND REBUILD ALL GEOMETRY
    # ==========================================================================
    print(f"[JOURNAL] STEP 2: Updating expressions in M1_Blank...")

    # Find and switch to M1_Blank part
    try:
        part3 = theSession.Parts.FindObject("M1_Blank")
        markId3 = theSession.SetUndoMark(
            NXOpen.Session.MarkVisibility.Visible, "Change Displayed Part"
        )
        status1, partLoadStatus3 = theSession.Parts.SetActiveDisplay(
            part3,
            NXOpen.DisplayPartOption.AllowAdditional,
            NXOpen.PartDisplayPartWorkPartOption.UseLast,
        )
        partLoadStatus3.Dispose()

        # Switch to modeling application for expression editing
        theSession.ApplicationSwitchImmediate("UG_APP_MODELING")

        workPart = theSession.Parts.Work

        # Create undo mark for expressions
        markId4 = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Visible, "Start")
        theSession.SetUndoMarkName(markId4, "Expressions Dialog")

        # Write expressions to a temp file and import (more reliable than editing one by one)
        exp_file_path = os.path.join(working_dir, "_temp_expressions.exp")
        with open(exp_file_path, "w") as f:
            for expr_name, expr_value in expression_updates.items():
                # Determine unit
                if "angle" in expr_name.lower() or "vertical" in expr_name.lower():
                    unit_str = "Degrees"
                else:
                    unit_str = "MilliMeter"
                f.write(f"[{unit_str}]{expr_name}={expr_value}\n")
                print(f"[JOURNAL]     {expr_name} = {expr_value} ({unit_str})")

        print(f"[JOURNAL]   Importing expressions from file...")
        markId_import = theSession.SetUndoMark(
            NXOpen.Session.MarkVisibility.Visible, "Import Expressions"
        )

        try:
            expModified, errorMessages = workPart.Expressions.ImportFromFile(
                exp_file_path, NXOpen.ExpressionCollection.ImportMode.Replace
            )
            print(f"[JOURNAL]   Expressions imported: {expModified} modified")
            if errorMessages:
                print(f"[JOURNAL]   Import errors: {errorMessages}")

            # Update geometry after import
            print(f"[JOURNAL]   Rebuilding M1_Blank geometry...")
            markId_update = theSession.SetUndoMark(
                NXOpen.Session.MarkVisibility.Invisible, "NX update"
            )
            nErrs = theSession.UpdateManager.DoUpdate(markId_update)
            theSession.DeleteUndoMark(markId_update, "NX update")
            print(f"[JOURNAL]   M1_Blank geometry rebuilt ({nErrs} errors)")

            # CRITICAL: Save M1_Blank after geometry update so FEM can read updated geometry
            print(f"[JOURNAL]   Saving M1_Blank...")
            partSaveStatus_blank = workPart.Save(
                NXOpen.BasePart.SaveComponents.TrueValue, NXOpen.BasePart.CloseAfterSave.FalseValue
            )
            partSaveStatus_blank.Dispose()
            print(f"[JOURNAL]   M1_Blank saved")

            # STEP 2a: EXTRACT MASS FROM M1_BLANK
            # Extract mass using MeasureManager after geometry is updated
            print(f"[JOURNAL]   Extracting mass from M1_Blank...")
            try:
                mass_kg = extract_part_mass(theSession, workPart, working_dir)
                print(f"[JOURNAL]   Mass extracted: {mass_kg:.6f} kg ({mass_kg * 1000:.2f} g)")
            except Exception as mass_err:
                print(f"[JOURNAL]   WARNING: Mass extraction failed: {mass_err}")

            updated_expressions = list(expression_updates.keys())

        except Exception as e:
            print(f"[JOURNAL]   ERROR importing expressions: {e}")
            updated_expressions = []

        # Clean up temp file
        try:
            os.remove(exp_file_path)
        except:
            pass

        theSession.SetUndoMarkName(markId4, "Expressions")

    except Exception as e:
        print(f"[JOURNAL]   ERROR updating expressions: {e}")

    # ==========================================================================
    # STEP 2b: UPDATE ALL LINKED GEOMETRY PARTS
    # ==========================================================================
    # CRITICAL: Must update ALL geometry parts that have linked expressions
    # before updating FEMs, otherwise interface nodes won't be coincident!
    print(f"[JOURNAL] STEP 2b: Updating all linked geometry parts...")

    # List of geometry parts that may have linked expressions from M1_Blank
    linked_geometry_parts = [
        "M1_Vertical_Support_Skeleton",
        # Add more parts here if the assembly has additional linked geometry
    ]

    for part_name in linked_geometry_parts:
        try:
            print(f"[JOURNAL]   Updating {part_name}...")
            linked_part = theSession.Parts.FindObject(part_name)

            markId_linked = theSession.SetUndoMark(
                NXOpen.Session.MarkVisibility.Visible, f"Update {part_name}"
            )
            status_linked, partLoadStatus_linked = theSession.Parts.SetActiveDisplay(
                linked_part,
                NXOpen.DisplayPartOption.AllowAdditional,
                NXOpen.PartDisplayPartWorkPartOption.UseLast,
            )
            partLoadStatus_linked.Dispose()

            # Switch to modeling application
            theSession.ApplicationSwitchImmediate("UG_APP_MODELING")

            # Update to propagate linked expression changes
            markId_linked_update = theSession.SetUndoMark(
                NXOpen.Session.MarkVisibility.Invisible, "NX update"
            )
            nErrs_linked = theSession.UpdateManager.DoUpdate(markId_linked_update)
            theSession.DeleteUndoMark(markId_linked_update, "NX update")
            print(f"[JOURNAL]     {part_name} geometry rebuilt ({nErrs_linked} errors)")

            # CRITICAL: Save part after geometry update so FEM can read updated geometry
            print(f"[JOURNAL]     Saving {part_name}...")
            partSaveStatus_linked = linked_part.Save(
                NXOpen.BasePart.SaveComponents.TrueValue, NXOpen.BasePart.CloseAfterSave.FalseValue
            )
            partSaveStatus_linked.Dispose()
            print(f"[JOURNAL]     {part_name} saved")

        except Exception as e:
            print(f"[JOURNAL]     WARNING: Could not update {part_name}: {e}")
            print(f"[JOURNAL]     (Part may not exist in this assembly)")

    # ==========================================================================
    # STEP 3: OPEN SIM AND UPDATE COMPONENT FEMs
    # ==========================================================================
    print(f"[JOURNAL] STEP 3: Opening sim and updating component FEMs...")

    # Try to find the sim part first (like the recorded journal does)
    # This ensures we're working with the same loaded sim part context
    sim_part_name = os.path.splitext(sim_filename)[0]  # e.g., "ASSY_M1_assyfem1_sim1"
    print(f"[JOURNAL]   Looking for sim part: {sim_part_name}")

    markId_sim = theSession.SetUndoMark(
        NXOpen.Session.MarkVisibility.Visible, "Change Displayed Part"
    )

    try:
        # First try to find it among loaded parts (like recorded journal)
        simPart1 = theSession.Parts.FindObject(sim_part_name)
        status_sim, partLoadStatus = theSession.Parts.SetActiveDisplay(
            simPart1,
            NXOpen.DisplayPartOption.AllowAdditional,
            NXOpen.PartDisplayPartWorkPartOption.UseLast,
        )
        partLoadStatus.Dispose()
        print(f"[JOURNAL]   Found and activated existing sim part")
    except:
        # Fallback: Open fresh if not found
        print(f"[JOURNAL]   Sim part not found, opening fresh: {sim_filename}")
        basePart, partLoadStatus = theSession.Parts.OpenActiveDisplay(
            sim_file_path, NXOpen.DisplayPartOption.AllowAdditional
        )
        partLoadStatus.Dispose()

    workSimPart = theSession.Parts.BaseWork
    displaySimPart = theSession.Parts.BaseDisplay
    theSession.ApplicationSwitchImmediate("UG_APP_SFEM")
    theSession.Post.UpdateUserGroupsFromSimPart(workSimPart)

    # Navigate component hierarchy
    try:
        rootComponent = workSimPart.ComponentAssembly.RootComponent
        component1 = rootComponent.FindObject("COMPONENT ASSY_M1_assyfem1 1")

        # Update M1_Blank_fem1
        print(f"[JOURNAL]   Updating M1_Blank_fem1...")
        try:
            component2 = component1.FindObject("COMPONENT M1_Blank_fem1 1")
            markId_fem1 = theSession.SetUndoMark(
                NXOpen.Session.MarkVisibility.Visible, "Make Work Part"
            )
            partLoadStatus5 = theSession.Parts.SetWorkComponent(
                component2,
                NXOpen.PartCollection.RefsetOption.Entire,
                NXOpen.PartCollection.WorkComponentOption.Visible,
            )
            workFemPart = theSession.Parts.BaseWork
            partLoadStatus5.Dispose()

            markId_update1 = theSession.SetUndoMark(
                NXOpen.Session.MarkVisibility.Visible, "Update FE Model"
            )
            fEModel1 = workFemPart.FindObject("FEModel")
            fEModel1.UpdateFemodel()
            print(f"[JOURNAL]     M1_Blank_fem1 updated")

            # CRITICAL: Save FEM file after update to persist mesh changes
            print(f"[JOURNAL]     Saving M1_Blank_fem1...")
            partSaveStatus_fem1 = workFemPart.Save(
                NXOpen.BasePart.SaveComponents.TrueValue, NXOpen.BasePart.CloseAfterSave.FalseValue
            )
            partSaveStatus_fem1.Dispose()
            print(f"[JOURNAL]     M1_Blank_fem1 saved")
        except Exception as e:
            print(f"[JOURNAL]     WARNING: M1_Blank_fem1: {e}")

        # Update M1_Vertical_Support_Skeleton_fem1
        print(f"[JOURNAL]   Updating M1_Vertical_Support_Skeleton_fem1...")
        try:
            component3 = component1.FindObject("COMPONENT M1_Vertical_Support_Skeleton_fem1 3")
            markId_fem2 = theSession.SetUndoMark(
                NXOpen.Session.MarkVisibility.Visible, "Make Work Part"
            )
            partLoadStatus6 = theSession.Parts.SetWorkComponent(
                component3,
                NXOpen.PartCollection.RefsetOption.Entire,
                NXOpen.PartCollection.WorkComponentOption.Visible,
            )
            workFemPart = theSession.Parts.BaseWork
            partLoadStatus6.Dispose()

            markId_update2 = theSession.SetUndoMark(
                NXOpen.Session.MarkVisibility.Visible, "Update FE Model"
            )
            fEModel2 = workFemPart.FindObject("FEModel")
            fEModel2.UpdateFemodel()
            print(f"[JOURNAL]     M1_Vertical_Support_Skeleton_fem1 updated")

            # CRITICAL: Save FEM file after update to persist mesh changes
            print(f"[JOURNAL]     Saving M1_Vertical_Support_Skeleton_fem1...")
            partSaveStatus_fem2 = workFemPart.Save(
                NXOpen.BasePart.SaveComponents.TrueValue, NXOpen.BasePart.CloseAfterSave.FalseValue
            )
            partSaveStatus_fem2.Dispose()
            print(f"[JOURNAL]     M1_Vertical_Support_Skeleton_fem1 saved")
        except Exception as e:
            print(f"[JOURNAL]     WARNING: M1_Vertical_Support_Skeleton_fem1: {e}")

    except Exception as e:
        print(f"[JOURNAL]   ERROR navigating component hierarchy: {e}")

    # ==========================================================================
    # STEP 4: MERGE DUPLICATE NODES
    # ==========================================================================
    print(f"[JOURNAL] STEP 4: Merging duplicate nodes...")

    try:
        # Switch to assembly FEM
        partLoadStatus8 = theSession.Parts.SetWorkComponent(
            component1,
            NXOpen.PartCollection.RefsetOption.Entire,
            NXOpen.PartCollection.WorkComponentOption.Visible,
        )
        workAssyFemPart = theSession.Parts.BaseWork
        displaySimPart = theSession.Parts.BaseDisplay
        partLoadStatus8.Dispose()
        print(f"[JOURNAL]   Switched to assembly FEM: {workAssyFemPart.Name}")

        markId_merge = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Visible, "Start")

        # WORKAROUND: Force display refresh before duplicate node check
        # The recorded journal does zoom operations before checking - this may
        # be needed to refresh the internal mesh representation
        try:
            displaySimPart.ModelingViews.WorkView.Fit()
            print(f"[JOURNAL]   Forced view Fit() to refresh display")
        except Exception as fit_err:
            print(f"[JOURNAL]   View Fit() failed (non-critical): {fit_err}")

        caePart1 = workAssyFemPart
        duplicateNodesCheckBuilder1 = caePart1.ModelCheckMgr.CreateDuplicateNodesCheckBuilder()

        # Set tolerance
        unit_tol = duplicateNodesCheckBuilder1.Tolerance.Units
        duplicateNodesCheckBuilder1.Tolerance.Units = unit_tol
        duplicateNodesCheckBuilder1.Tolerance.SetFormula("0.01")
        print(f"[JOURNAL]   Tolerance: 0.01 mm")

        # Enable occurrence node merge - CRITICAL for assembly FEM
        duplicateNodesCheckBuilder1.MergeOccurrenceNodes = True
        print(f"[JOURNAL]   MergeOccurrenceNodes: True")

        theSession.SetUndoMarkName(markId_merge, "Duplicate Nodes Dialog")

        # Configure display settings
        displaysettings1 = NXOpen.CAE.ModelCheck.DuplicateNodesCheckBuilder.DisplaySettings()
        displaysettings1.ShowDuplicateNodes = True
        displaysettings1.ShowMergedNodeLabels = False
        displaysettings1.ShowRetainedNodeLabels = False
        displaysettings1.KeepNodesColor = displaySimPart.Colors.Find("Blue")
        displaysettings1.MergeNodesColor = displaySimPart.Colors.Find("Yellow")
        displaysettings1.UnableToMergeNodesColor = displaySimPart.Colors.Find("Red")
        duplicateNodesCheckBuilder1.DisplaySettingsData = displaysettings1

        # Check scope
        duplicateNodesCheckBuilder1.CheckScopeOption = NXOpen.CAE.ModelCheck.CheckScope.Displayed
        print(f"[JOURNAL]   CheckScope: Displayed")

        # Identify duplicates
        print(f"[JOURNAL]   Identifying duplicate nodes...")
        numDuplicates = duplicateNodesCheckBuilder1.IdentifyDuplicateNodes()
        print(f"[JOURNAL]   IdentifyDuplicateNodes returned: {numDuplicates}")

        # WORKAROUND: In batch mode, IdentifyDuplicateNodes() often returns None
        # even when duplicates exist. The recorded NX journal doesn't check the
        # return value - it just calls MergeDuplicateNodes unconditionally.
        # So we do the same: always attempt to merge.
        print(f"[JOURNAL]   Attempting to merge duplicate nodes...")
        try:
            numMerged = duplicateNodesCheckBuilder1.MergeDuplicateNodes()
            print(f"[JOURNAL]   MergeDuplicateNodes returned: {numMerged}")
            if numMerged is not None and numMerged > 0:
                print(f"[JOURNAL]   Successfully merged {numMerged} duplicate node sets")
            elif numMerged == 0:
                print(f"[JOURNAL]   No nodes were merged (0 returned)")
                if numDuplicates is None:
                    print(
                        f"[JOURNAL]   WARNING: IdentifyDuplicateNodes returned None - mesh may need display refresh"
                    )
            else:
                print(f"[JOURNAL]   MergeDuplicateNodes returned None - batch mode limitation")
        except Exception as merge_error:
            print(f"[JOURNAL]   MergeDuplicateNodes failed: {merge_error}")
            if numDuplicates is None:
                print(
                    f"[JOURNAL]   This combined with IdentifyDuplicateNodes=None suggests display issue"
                )

        theSession.SetUndoMarkName(markId_merge, "Duplicate Nodes")
        duplicateNodesCheckBuilder1.Destroy()
        theSession.DeleteUndoMark(markId_merge, None)

    except Exception as e:
        print(f"[JOURNAL]   WARNING: Node merge: {e}")
        import traceback

        traceback.print_exc()

    # ==========================================================================
    # STEP 5: RESOLVE LABEL CONFLICTS
    # ==========================================================================
    print(f"[JOURNAL] STEP 5: Resolving label conflicts...")

    try:
        markId_labels = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Visible, "Start")

        assyFemPart1 = workAssyFemPart
        assemblyLabelManagerBuilder1 = assyFemPart1.CreateAssemblyLabelManagerBuilder()

        theSession.SetUndoMarkName(markId_labels, "Assembly Label Manager Dialog")

        markId_labels2 = theSession.SetUndoMark(
            NXOpen.Session.MarkVisibility.Invisible, "Assembly Label Manager"
        )

        # Set offsets for each FE model occurrence
        # These offsets ensure unique node/element labels across components
        entitytypes = [
            NXOpen.CAE.AssemblyLabelManagerBuilder.EntityType.Node,
            NXOpen.CAE.AssemblyLabelManagerBuilder.EntityType.Element,
            NXOpen.CAE.AssemblyLabelManagerBuilder.EntityType.Csys,
            NXOpen.CAE.AssemblyLabelManagerBuilder.EntityType.Physical,
            NXOpen.CAE.AssemblyLabelManagerBuilder.EntityType.Group,
            NXOpen.CAE.AssemblyLabelManagerBuilder.EntityType.Ply,
            NXOpen.CAE.AssemblyLabelManagerBuilder.EntityType.Ssmo,
        ]

        # Apply offsets to each occurrence (values from recorded journal)
        occurrence_offsets = [
            ("FEModelOccurrence[3]", 2),
            ("FEModelOccurrence[4]", 74),
            ("FEModelOccurrence[5]", 146),
            ("FEModelOccurrence[7]", 218),
        ]

        for occ_name, offset_val in occurrence_offsets:
            try:
                fEModelOcc = workAssyFemPart.FindObject(occ_name)
                offsets = [offset_val] * 7
                assemblyLabelManagerBuilder1.SetFEModelOccOffsets(fEModelOcc, entitytypes, offsets)
            except:
                pass  # Some occurrences may not exist

        nXObject1 = assemblyLabelManagerBuilder1.Commit()

        theSession.DeleteUndoMark(markId_labels2, None)
        theSession.SetUndoMarkName(markId_labels, "Assembly Label Manager")
        assemblyLabelManagerBuilder1.Destroy()

        print(f"[JOURNAL]   Label conflicts resolved")

    except Exception as e:
        print(f"[JOURNAL]   WARNING: Label management: {e}")

    # ==========================================================================
    # STEP 5b: SAVE ASSEMBLY FEM
    # ==========================================================================
    print(f"[JOURNAL] STEP 5b: Saving assembly FEM after all updates...")
    try:
        # Save the assembly FEM to persist all mesh updates and node merges
        partSaveStatus_afem = workAssyFemPart.Save(
            NXOpen.BasePart.SaveComponents.TrueValue, NXOpen.BasePart.CloseAfterSave.FalseValue
        )
        partSaveStatus_afem.Dispose()
        print(f"[JOURNAL]   Assembly FEM saved: {workAssyFemPart.Name}")
    except Exception as e:
        print(f"[JOURNAL]   WARNING: Could not save assembly FEM: {e}")

    # ==========================================================================
    # STEP 6: SOLVE
    # ==========================================================================
    print(f"[JOURNAL] STEP 6: Solving simulation...")

    try:
        # Return to sim level by setting null component
        partLoadStatus9 = theSession.Parts.SetWorkComponent(
            NXOpen.Assemblies.Component.Null,
            NXOpen.PartCollection.RefsetOption.Entire,
            NXOpen.PartCollection.WorkComponentOption.Visible,
        )
        workSimPart = theSession.Parts.BaseWork
        partLoadStatus9.Dispose()

        # Set up solve
        markId_solve = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Visible, "Start")
        theSession.SetUndoMarkName(markId_solve, "Solve Dialog")

        markId_solve2 = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Invisible, "Solve")

        theCAESimSolveManager = NXOpen.CAE.SimSolveManager.GetSimSolveManager(theSession)

        simSimulation1 = workSimPart.FindObject("Simulation")
        sol_name = solution_name if solution_name else "Solution 1"
        simSolution1 = simSimulation1.FindObject(f"Solution[{sol_name}]")

        psolutions1 = [simSolution1]

        print(f"[JOURNAL]   Solving: {sol_name} (Foreground mode)")
        numsolved, numfailed, numskipped = theCAESimSolveManager.SolveChainOfSolutions(
            psolutions1,
            NXOpen.CAE.SimSolution.SolveOption.Solve,
            NXOpen.CAE.SimSolution.SetupCheckOption.CompleteCheckAndOutputErrors,
            NXOpen.CAE.SimSolution.SolveMode.Foreground,  # Use Foreground to ensure OP2 is complete
        )

        theSession.DeleteUndoMark(markId_solve2, None)
        theSession.SetUndoMarkName(markId_solve, "Solve")

        print(
            f"[JOURNAL] Solve completed: {numsolved} solved, {numfailed} failed, {numskipped} skipped"
        )

        # ==========================================================================
        # STEP 7: SAVE ALL - Save all modified parts (FEM, SIM, PRT)
        # ==========================================================================
        print(f"[JOURNAL] STEP 7: Saving all modified parts...")
        try:
            anyPartsModified, partSaveStatus_all = theSession.Parts.SaveAll()
            partSaveStatus_all.Dispose()
            print(f"[JOURNAL]   SaveAll completed (parts modified: {anyPartsModified})")
        except Exception as e:
            print(f"[JOURNAL]   WARNING: SaveAll failed: {e}")

        return numfailed == 0

    except Exception as e:
        print(f"[JOURNAL] ERROR solving: {e}")
        import traceback

        traceback.print_exc()
        return False


def solve_simple_workflow(
    theSession, sim_file_path, solution_name, expression_updates, working_dir
):
    """
    Workflow for single-part simulations with optional expression updates.

    For single-part FEMs (Bracket.prt -> Bracket_fem1.fem -> Bracket_sim1.sim):
    1. Open the .sim file (this loads .fem and .prt)
    2. If expression_updates: find the geometry .prt, update expressions, rebuild
    3. Update the FEM mesh
    4. Solve
    """
    print(f"[JOURNAL] Opening simulation: {sim_file_path}")

    # Open the .sim file
    print(f"[JOURNAL]   sim_file_path = {sim_file_path}")
    print(f"[JOURNAL]   File exists: {os.path.exists(sim_file_path)}")
    basePart1, partLoadStatus1 = theSession.Parts.OpenActiveDisplay(
        sim_file_path, NXOpen.DisplayPartOption.AllowAdditional
    )

    # Diagnostic: check load status
    print(f"[JOURNAL]   OpenActiveDisplay result: basePart1={basePart1}")
    print(f"[JOURNAL]   basePart1 type: {type(basePart1).__name__}" if basePart1 else "[JOURNAL]   basePart1 is None!")
    if basePart1:
        print(f"[JOURNAL]   basePart1.Name: {basePart1.Name}")
        print(f"[JOURNAL]   basePart1.FullPath: {basePart1.FullPath}")
    try:
        n_statuses = partLoadStatus1.NumberUnloadedParts
        print(f"[JOURNAL]   Unloaded parts: {n_statuses}")
        for i in range(n_statuses):
            name = partLoadStatus1.GetPartName(i)
            status = partLoadStatus1.GetStatus(i)
            status_desc = partLoadStatus1.GetStatusDescription(i)
            print(f"[JOURNAL]     Part[{i}]: {name} — status={status} ({status_desc})")
    except Exception as e:
        print(f"[JOURNAL]   Could not read load status details: {e}")
    partLoadStatus1.Dispose()

    workSimPart = theSession.Parts.BaseWork
    print(f"[JOURNAL]   BaseWork: {workSimPart.Name if workSimPart else 'None'}")
    print(f"[JOURNAL]   Parts count: {sum(1 for _ in theSession.Parts)}")

    # =========================================================================
    # STEP 1: UPDATE EXPRESSIONS IN GEOMETRY PART (if any)
    # =========================================================================
    if expression_updates:
        print(f"[JOURNAL] STEP 1: Updating expressions in geometry part...")

        # List all loaded parts for debugging
        print(f"[JOURNAL]   Currently loaded parts:")
        for part in theSession.Parts:
            print(f"[JOURNAL]     - {part.Name} (type: {type(part).__name__})")

        # NX doesn't automatically load the geometry .prt when opening a SIM file
        # We need to find and load it explicitly from the working directory
        geom_part = None

        # First, try to find an already loaded geometry part
        for part in theSession.Parts:
            part_name = part.Name.lower()
            part_type = type(part).__name__

            # Skip FEM and SIM parts by type
            if "fem" in part_type.lower() or "sim" in part_type.lower():
                continue

            # Skip parts with _fem or _sim in name
            if "_fem" in part_name or "_sim" in part_name:
                continue

            geom_part = part
            print(f"[JOURNAL]   Found geometry part (already loaded): {part.Name}")
            break

        # If not found, try to load the geometry .prt file from working directory
        if geom_part is None:
            print(f"[JOURNAL]   Geometry part not loaded, searching for .prt file...")
            for filename in os.listdir(working_dir):
                # Skip idealized parts (_i.prt), FEM parts, and SIM parts
                if (
                    filename.endswith(".prt")
                    and "_fem" not in filename.lower()
                    and "_sim" not in filename.lower()
                    and "_i.prt" not in filename.lower()
                ):
                    prt_path = os.path.join(working_dir, filename)
                    print(f"[JOURNAL]   Loading geometry part: {filename}")
                    print(f"[JOURNAL]     Full path: {prt_path}")
                    print(f"[JOURNAL]     File exists: {os.path.exists(prt_path)}")
                    print(f"[JOURNAL]     File size: {os.path.getsize(prt_path) if os.path.exists(prt_path) else 'N/A'}")
                    try:
                        loaded_part, partLoadStatus = theSession.Parts.Open(prt_path)
                        try:
                            n_unloaded = partLoadStatus.NumberUnloadedParts
                            if n_unloaded > 0:
                                print(f"[JOURNAL]     Parts.Open unloaded parts: {n_unloaded}")
                                for i in range(n_unloaded):
                                    pn = partLoadStatus.GetPartName(i)
                                    ps = partLoadStatus.GetStatusDescription(i)
                                    print(f"[JOURNAL]       [{i}]: {pn} — {ps}")
                        except:
                            pass
                        partLoadStatus.Dispose()
                        # Check if load actually succeeded (Parts.Open can return None)
                        if loaded_part is not None:
                            geom_part = loaded_part
                            print(f"[JOURNAL]   Geometry part loaded: {geom_part.Name}")
                            break
                        else:
                            print(f"[JOURNAL]   WARNING: Parts.Open returned None for {filename}")
                            # Check if part got loaded anyway under a different reference
                            print(f"[JOURNAL]   Parts after Open attempt: {sum(1 for _ in theSession.Parts)}")
                            for p in theSession.Parts:
                                print(f"[JOURNAL]     - {p.Name} ({type(p).__name__})")
                    except Exception as e:
                        print(f"[JOURNAL]   WARNING: Could not load {filename}: {e}")
                        import traceback
                        traceback.print_exc()

        if geom_part:
            try:
                # Switch to the geometry part for expression editing
                markId_expr = theSession.SetUndoMark(
                    NXOpen.Session.MarkVisibility.Visible, "Update Expressions"
                )
                status, partLoadStatus = theSession.Parts.SetActiveDisplay(
                    geom_part,
                    NXOpen.DisplayPartOption.AllowAdditional,
                    NXOpen.PartDisplayPartWorkPartOption.UseLast,
                )
                partLoadStatus.Dispose()

                # Switch to modeling application for expression editing
                theSession.ApplicationSwitchImmediate("UG_APP_MODELING")

                workPart = theSession.Parts.Work

                # Write expressions to temp file and import
                exp_file_path = os.path.join(working_dir, "_temp_expressions.exp")

                # Known integer/constant expressions (no unit)
                CONSTANT_EXPRESSIONS = {
                    "hole_count",
                }

                with open(exp_file_path, "w") as f:
                    for expr_name, expr_value in expression_updates.items():
                        # Determine unit based on expression type
                        if expr_name in CONSTANT_EXPRESSIONS:
                            unit_str = "Constant"
                            # Write as integer if it's a whole number
                            if expr_value == int(expr_value):
                                expr_value = int(expr_value)
                        elif "angle" in expr_name.lower():
                            unit_str = "Degrees"
                        else:
                            unit_str = "MilliMeter"
                        f.write(f"[{unit_str}]{expr_name}={expr_value}\n")
                        print(f"[JOURNAL]     {expr_name} = {expr_value} ({unit_str})")

                print(f"[JOURNAL]   Importing expressions...")
                expModified, errorMessages = workPart.Expressions.ImportFromFile(
                    exp_file_path, NXOpen.ExpressionCollection.ImportMode.Replace
                )
                print(f"[JOURNAL]   Expressions modified: {expModified}")
                if errorMessages:
                    print(f"[JOURNAL]   Import messages: {errorMessages}")

                # Update geometry
                print(f"[JOURNAL]   Rebuilding geometry...")
                markId_update = theSession.SetUndoMark(
                    NXOpen.Session.MarkVisibility.Invisible, "NX update"
                )
                nErrs = theSession.UpdateManager.DoUpdate(markId_update)
                theSession.DeleteUndoMark(markId_update, "NX update")
                print(f"[JOURNAL]   Geometry rebuilt ({nErrs} errors)")

                # Save geometry part
                print(f"[JOURNAL]   Saving geometry part...")
                partSaveStatus_geom = workPart.Save(
                    NXOpen.BasePart.SaveComponents.TrueValue,
                    NXOpen.BasePart.CloseAfterSave.FalseValue,
                )
                partSaveStatus_geom.Dispose()

                # Clean up temp file
                try:
                    os.remove(exp_file_path)
                except:
                    pass

            except Exception as e:
                print(f"[JOURNAL]   ERROR updating expressions: {e}")
                import traceback

                traceback.print_exc()
        else:
            print(f"[JOURNAL]   WARNING: Could not find geometry part for expression updates!")

    # =========================================================================
    # STEP 2: UPDATE FEM MESH (if expressions were updated)
    # =========================================================================
    if expression_updates:
        print(f"[JOURNAL] STEP 2: Updating FEM mesh...")

        # First, load the idealized part if it exists (required for mesh update chain)
        # The chain is: .prt (geometry) -> _i.prt (idealized) -> .fem (mesh)
        idealized_part = None
        for filename in os.listdir(working_dir):
            if "_i.prt" in filename.lower():
                idealized_path = os.path.join(working_dir, filename)
                print(f"[JOURNAL]   Loading idealized part: {filename}")
                try:
                    loaded_part, partLoadStatus = theSession.Parts.Open(idealized_path)
                    partLoadStatus.Dispose()
                    # Check if load actually succeeded (Parts.Open can return None)
                    if loaded_part is not None:
                        idealized_part = loaded_part
                        print(f"[JOURNAL]   Idealized part loaded: {idealized_part.Name}")
                    else:
                        print(f"[JOURNAL]   WARNING: Parts.Open returned None for idealized part")
                except Exception as e:
                    print(f"[JOURNAL]   WARNING: Could not load idealized part: {e}")
                break

        # Find the FEM part (must be .fem, NOT _i.prt idealized part)
        fem_part = None
        for part in theSession.Parts:
            part_name = part.Name.lower()
            # Match FEM parts but exclude idealized parts (_i)
            if ("_fem" in part_name or part_name.endswith(".fem")) and "_i" not in part_name.split("_fem")[-1]:
                fem_part = part
                print(f"[JOURNAL]   Found FEM part: {part.Name}")
                break

        # If not found by name, try loading .fem file from working directory
        if fem_part is None:
            for filename in os.listdir(working_dir):
                if filename.lower().endswith(".fem"):
                    fem_path = os.path.join(working_dir, filename)
                    print(f"[JOURNAL]   Loading FEM file: {filename}")
                    try:
                        loaded_part, partLoadStatus = theSession.Parts.Open(fem_path)
                        partLoadStatus.Dispose()
                        if loaded_part is not None:
                            fem_part = loaded_part
                            print(f"[JOURNAL]   FEM part loaded: {fem_part.Name}")
                        break
                    except Exception as e:
                        print(f"[JOURNAL]   WARNING: Could not load FEM: {e}")

        if fem_part:
            try:
                # Switch to FEM part - CRITICAL: Use SameAsDisplay to make FEM the work part
                # This is required for UpdateFemodel() to properly regenerate the mesh
                # Reference: tests/journal_with_regenerate.py line 76
                print(f"[JOURNAL]   Switching to FEM part: {fem_part.Name}")
                status, partLoadStatus = theSession.Parts.SetActiveDisplay(
                    fem_part,
                    NXOpen.DisplayPartOption.AllowAdditional,
                    NXOpen.PartDisplayPartWorkPartOption.SameAsDisplay,  # Critical fix!
                )
                partLoadStatus.Dispose()

                # Switch to FEM application
                theSession.ApplicationSwitchImmediate("UG_APP_SFEM")

                # Update the FE model
                workFemPart = theSession.Parts.BaseWork
                feModel = workFemPart.FindObject("FEModel")

                print(f"[JOURNAL]   Updating FE model...")
                feModel.UpdateFemodel()
                print(f"[JOURNAL]   FE model updated")

                # Save FEM
                partSaveStatus_fem = workFemPart.Save(
                    NXOpen.BasePart.SaveComponents.TrueValue,
                    NXOpen.BasePart.CloseAfterSave.FalseValue,
                )
                partSaveStatus_fem.Dispose()
                print(f"[JOURNAL]   FEM saved")

            except Exception as e:
                print(f"[JOURNAL]   ERROR updating FEM: {e}")
                import traceback

                traceback.print_exc()

    # =========================================================================
    # STEP 3: SWITCH BACK TO SIM AND SOLVE
    # =========================================================================
    print(f"[JOURNAL] STEP 3: Solving simulation...")

    # Switch back to sim part
    status, partLoadStatus = theSession.Parts.SetActiveDisplay(
        workSimPart,
        NXOpen.DisplayPartOption.AllowAdditional,
        NXOpen.PartDisplayPartWorkPartOption.UseLast,
    )
    partLoadStatus.Dispose()

    theSession.ApplicationSwitchImmediate("UG_APP_SFEM")
    theSession.Post.UpdateUserGroupsFromSimPart(workSimPart)

    # Set up solve
    markId_solve = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Visible, "Start")
    theSession.SetUndoMarkName(markId_solve, "Solve Dialog")

    markId_solve2 = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Invisible, "Solve")

    theCAESimSolveManager = NXOpen.CAE.SimSolveManager.GetSimSolveManager(theSession)

    simSimulation1 = workSimPart.FindObject("Simulation")
    sol_name = solution_name if solution_name else "Solution 1"
    simSolution1 = simSimulation1.FindObject(f"Solution[{sol_name}]")

    psolutions1 = [simSolution1]

    print(f"[JOURNAL]   Solving: {sol_name}")
    numsolved, numfailed, numskipped = theCAESimSolveManager.SolveChainOfSolutions(
        psolutions1,
        NXOpen.CAE.SimSolution.SolveOption.Solve,
        NXOpen.CAE.SimSolution.SetupCheckOption.CompleteCheckAndOutputErrors,
        NXOpen.CAE.SimSolution.SolveMode.Foreground,  # Use Foreground to wait for completion
    )

    theSession.DeleteUndoMark(markId_solve2, None)
    theSession.SetUndoMarkName(markId_solve, "Solve")

    print(
        f"[JOURNAL] Solve completed: {numsolved} solved, {numfailed} failed, {numskipped} skipped"
    )

    # Extract mass from geometry part expression (p173) and write to temp file
    try:
        mass_value = None
        # Find geometry part (Beam.prt)
        for part in theSession.Parts:
            part_type = type(part).__name__
            if "fem" not in part_type.lower() and "sim" not in part_type.lower():
                # This is the geometry part — look for mass expression
                for expr in part.Expressions:
                    if expr.Name == "p173":
                        mass_value = expr.Value
                        print(f"[JOURNAL] Mass expression p173 = {mass_value}")
                        break
                break

        if mass_value is not None:
            mass_file = os.path.join(working_dir, "_temp_mass.txt")
            with open(mass_file, "w") as f:
                f.write(f"p173={mass_value}\n")
            print(f"[JOURNAL] Wrote mass to {mass_file}")
        else:
            print(f"[JOURNAL] WARNING: Could not find mass expression p173")
    except Exception as e:
        print(f"[JOURNAL] WARNING: Mass extraction failed: {e}")

    # Save all
    try:
        anyPartsModified, partSaveStatus = theSession.Parts.SaveAll()
        partSaveStatus.Dispose()
        print(f"[JOURNAL] Saved all parts!")
    except:
        pass

    return numfailed == 0


if __name__ == "__main__":
    success = main(sys.argv[1:])
    sys.exit(0 if success else 1)