"""
NX Journal: Intelligent Model Discovery
========================================

This journal scans a .sim file and reports:
- All solutions (names, types, solver types)
- All expressions in the model (potential design variables)
- FEM properties (mesh info, materials)
- Output file locations

Usage:
    run_journal.exe discover_model.py -- sim_file.sim

Output:
    JSON to stdout with all discovered information
"""

import sys
import json
import NXOpen
import NXOpen.CAE

def discover_model(sim_path):
    """Discover all information about an NX simulation model."""

    result = {
        'sim_file': sim_path,
        'success': False,
        'solutions': [],
        'expressions': [],
        'materials': [],
        'mesh_info': {},
        'linked_parts': [],
        'errors': []
    }

    try:
        theSession = NXOpen.Session.GetSession()

        # Set load options to load from the directory containing the sim file
        import os
        sim_dir = os.path.dirname(os.path.abspath(sim_path))
        sim_name = os.path.basename(sim_path)

        loadOptions = theSession.Parts.LoadOptions
        loadOptions.ComponentLoadMethod = NXOpen.LoadOptions.LoadComponents.AllComponents
        loadOptions.SetSearchDirectories([sim_dir])

        print(f"[DISCOVER] Opening simulation: {sim_path}", file=sys.stderr)

        # Open the simulation file
        basePart, partLoadStatus = theSession.Parts.OpenActiveDisplay(sim_path, NXOpen.DisplayPartOption.ReplaceExisting)
        partLoadStatus.Dispose()

        if basePart is None:
            result['errors'].append("Failed to open simulation file")
            print(json.dumps(result))
            return

        # Get the simulation part
        simPart = theSession.Parts.BaseWork

        # ============================================================
        # DISCOVER SOLUTIONS
        # ============================================================
        print(f"[DISCOVER] Scanning for solutions...", file=sys.stderr)

        try:
            simSimulation = simPart.Simulation

            # Try to find solutions by iterating through common naming patterns
            solution_patterns = [
                "Solution 1", "Solution 2", "Solution 3", "Solution 4", "Solution 5",
                "solution_1", "solution_2", "solution_3",
                "Static", "Modal", "Buckling", "Frequency Response",
                "Thermal", "Heat Transfer", "Transient",
            ]

            found_solutions = []

            for pattern in solution_patterns:
                try:
                    sol_obj_name = f"Solution[{pattern}]"
                    sol = simSimulation.FindObject(sol_obj_name)
                    if sol:
                        sol_info = {
                            'name': pattern,
                            'object_name': sol_obj_name,
                            'type': str(type(sol).__name__),
                        }

                        # Try to get solver type
                        try:
                            sol_info['solver'] = str(sol.SolutionType)
                        except:
                            pass

                        # Try to get analysis type
                        try:
                            sol_info['analysis_type'] = str(sol.AnalysisType) if hasattr(sol, 'AnalysisType') else 'Unknown'
                        except:
                            pass

                        found_solutions.append(sol_info)
                        print(f"[DISCOVER]   Found solution: {pattern}", file=sys.stderr)
                except:
                    pass

            # Also try iterating by number (Solution 1, Solution 2, etc.)
            for i in range(1, 20):
                pattern = f"Solution {i}"
                if any(s['name'] == pattern for s in found_solutions):
                    continue
                try:
                    sol_obj_name = f"Solution[{pattern}]"
                    sol = simSimulation.FindObject(sol_obj_name)
                    if sol:
                        sol_info = {
                            'name': pattern,
                            'object_name': sol_obj_name,
                            'type': str(type(sol).__name__),
                        }
                        found_solutions.append(sol_info)
                        print(f"[DISCOVER]   Found solution: {pattern}", file=sys.stderr)
                except:
                    break  # Stop when we hit missing numbers

            result['solutions'] = found_solutions
            print(f"[DISCOVER] Found {len(found_solutions)} solution(s)", file=sys.stderr)

        except Exception as e:
            result['errors'].append(f"Error scanning solutions: {str(e)}")

        # ============================================================
        # DISCOVER EXPRESSIONS (Potential Design Variables)
        # ============================================================
        print(f"[DISCOVER] Scanning for expressions...", file=sys.stderr)

        try:
            # Find all linked parts (geometry parts)
            linked_parts = []

            # Check for common part naming patterns based on sim file name
            sim_base = os.path.splitext(sim_name)[0]  # e.g., "Bracket_sim1"

            # Extract the base model name (remove _sim1, _fem1, etc.)
            import re
            base_match = re.match(r'(.+?)(?:_sim\d*|_fem\d*)?$', sim_base, re.IGNORECASE)
            if base_match:
                model_base = base_match.group(1)  # e.g., "Bracket"
            else:
                model_base = sim_base

            print(f"[DISCOVER]   Model base name: {model_base}", file=sys.stderr)

            # Look for the geometry part
            for part in theSession.Parts:
                part_name = part.Name if hasattr(part, 'Name') else str(part)

                # Check if this is a geometry part (not sim or fem)
                if model_base.lower() in part_name.lower():
                    is_sim = '_sim' in part_name.lower() or part_name.lower().endswith('.sim')
                    is_fem = '_fem' in part_name.lower() or part_name.lower().endswith('.fem')

                    if not is_sim and not is_fem:
                        linked_parts.append(part_name)
                        print(f"[DISCOVER]   Found geometry part: {part_name}", file=sys.stderr)

                        # Get expressions from this part
                        try:
                            for expr in part.Expressions:
                                expr_info = {
                                    'name': expr.Name,
                                    'value': expr.Value if hasattr(expr, 'Value') else None,
                                    'formula': expr.RightHandSide if hasattr(expr, 'RightHandSide') else None,
                                    'type': str(expr.Type) if hasattr(expr, 'Type') else 'Unknown',
                                    'units': str(expr.Units) if hasattr(expr, 'Units') else None,
                                    'part': part_name
                                }

                                # Filter out system expressions
                                if not expr.Name.startswith('p') or not expr.Name[1:].isdigit():
                                    result['expressions'].append(expr_info)

                        except Exception as e:
                            result['errors'].append(f"Error reading expressions from {part_name}: {str(e)}")

            result['linked_parts'] = linked_parts
            print(f"[DISCOVER] Found {len(result['expressions'])} expression(s)", file=sys.stderr)

        except Exception as e:
            result['errors'].append(f"Error scanning expressions: {str(e)}")

        # ============================================================
        # DISCOVER MESH INFO
        # ============================================================
        print(f"[DISCOVER] Scanning for mesh info...", file=sys.stderr)

        try:
            # Look for FEM parts
            for part in theSession.Parts:
                part_name = part.Name if hasattr(part, 'Name') else str(part)

                if '_fem' in part_name.lower() or part_name.lower().endswith('.fem'):
                    print(f"[DISCOVER]   Found FEM part: {part_name}", file=sys.stderr)

                    mesh_info = {
                        'fem_part': part_name,
                        'element_count': 0,
                        'node_count': 0
                    }

                    # Try to get mesh statistics
                    try:
                        if hasattr(part, 'FEModel'):
                            feModel = part.FEModel
                            if feModel:
                                # Get element count
                                try:
                                    mesh_info['element_count'] = feModel.FeelementLabelMap.Size
                                except:
                                    pass

                                # Get node count
                                try:
                                    mesh_info['node_count'] = feModel.FenodeLabelMap.Size
                                except:
                                    pass
                    except:
                        pass

                    result['mesh_info'] = mesh_info
                    break

        except Exception as e:
            result['errors'].append(f"Error scanning mesh: {str(e)}")

        result['success'] = True
        print(f"[DISCOVER] Discovery complete!", file=sys.stderr)

    except Exception as e:
        result['errors'].append(f"Fatal error: {str(e)}")
        import traceback
        traceback.print_exc(file=sys.stderr)

    # Output JSON result to stdout
    print(json.dumps(result, indent=2))
    return result


if __name__ == '__main__':
    if len(sys.argv) < 2:
        print(json.dumps({'success': False, 'error': 'No sim file specified'}))
        sys.exit(1)

    sim_path = sys.argv[1]
    discover_model(sim_path)
