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

    # Get working directory
    working_dir = os.path.dirname(os.path.abspath(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
    basePart1, partLoadStatus1 = theSession.Parts.OpenActiveDisplay(
        sim_file_path,
        NXOpen.DisplayPartOption.AllowAdditional
    )
    partLoadStatus1.Dispose()

    workSimPart = theSession.Parts.BaseWork

    # =========================================================================
    # 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):
                if filename.endswith('.prt') and '_fem' not in filename.lower() and '_sim' not in filename.lower():
                    prt_path = os.path.join(working_dir, filename)
                    print(f"[JOURNAL]   Loading geometry part: {filename}")
                    try:
                        geom_part, partLoadStatus = theSession.Parts.Open(prt_path)
                        partLoadStatus.Dispose()
                        print(f"[JOURNAL]   Geometry part loaded: {geom_part.Name}")
                        break
                    except Exception as e:
                        print(f"[JOURNAL]   WARNING: Could not load {filename}: {e}")

        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")
                with open(exp_file_path, 'w') as f:
                    for expr_name, expr_value in expression_updates.items():
                        # Determine unit based on name
                        if '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:
                    idealized_part, partLoadStatus = theSession.Parts.Open(idealized_path)
                    partLoadStatus.Dispose()
                    print(f"[JOURNAL]   Idealized part loaded: {idealized_part.Name}")
                except Exception as e:
                    print(f"[JOURNAL]   WARNING: Could not load idealized part: {e}")
                break

        # Find the FEM part
        fem_part = None
        for part in theSession.Parts:
            if '_fem' in part.Name.lower() or part.Name.lower().endswith('.fem'):
                fem_part = part
                print(f"[JOURNAL]   Found FEM part: {part.Name}")
                break

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

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