Atomizer/optimization_engine/nx_updater.py

"""
NX Parameter Updater

Updates design variable values in NX .prt files.

This module can read expressions in two ways:
1. Parse .exp files (NX native export format) - RECOMMENDED, captures ALL expressions
2. Parse binary .prt files directly - fallback method, may miss some expressions

For updating values:
1. Binary .prt file modification (current implementation)
2. Future: Use NXOpen API if NX is running

The .exp format is preferred for reading because it captures:
- All expression types (formulas, references, constants)
- Unitless expressions
- Complete accuracy
"""

from pathlib import Path
from typing import Dict, List, Optional
import re
import shutil
import subprocess
from datetime import datetime
import sys

# Import centralized configuration
sys.path.insert(0, str(Path(__file__).parent.parent))
from config import NX_RUN_JOURNAL


class NXParameterUpdater:
    """
    Updates parametric expression values in NX .prt files.

    NX Expression Formats:

    Binary .prt format:
        #(Number [mm]) tip_thickness: 20.0;
        *(Number [degrees]) support_angle: 30.0;

    .exp export format (RECOMMENDED for reading):
        [MilliMeter]beam_length=5000
        [Kilogram]mass=973.968443678471
        hole_count=10
        Pattern_p7=hole_count
    """

    def __init__(self, prt_file_path: Path, backup: bool = True, nx_run_journal_path: Optional[Path] = None):
        """
        Initialize updater for a specific .prt file.

        Args:
            prt_file_path: Path to NX .prt file
            backup: If True, create backup before modifying
            nx_run_journal_path: Path to NX run_journal.exe (for .exp export)
                               If None, uses default NX 2412 path
        """
        self.prt_path = Path(prt_file_path)

        if not self.prt_path.exists():
            raise FileNotFoundError(f".prt file not found: {prt_file_path}")

        self.backup_enabled = backup
        self.content = None
        self.text_content = None

        # Default NX run_journal.exe path
        if nx_run_journal_path is None:
            self.nx_run_journal_path = NX_RUN_JOURNAL
        else:
            self.nx_run_journal_path = Path(nx_run_journal_path)

        self._load_file()

    def _load_file(self):
        """Load .prt file as binary."""
        with open(self.prt_path, 'rb') as f:
            self.content = bytearray(f.read())

        # Decode as latin-1 for text operations (preserves all bytes)
        self.text_content = self.content.decode('latin-1', errors='ignore')

    def _create_backup(self):
        """Create timestamped backup of original file."""
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        backup_path = self.prt_path.with_suffix(f'.prt.bak_{timestamp}')
        shutil.copy2(self.prt_path, backup_path)
        print(f"Backup created: {backup_path}")
        return backup_path

    def find_expressions(self) -> List[Dict[str, any]]:
        """
        Find all expressions in the .prt file.

        Returns:
            List of dicts with name, value, units
        """
        expressions = []

        # Pattern for NX expressions (with optional units):
        # #(Number [mm]) tip_thickness: 20.0;       - with units
        # *(Number [mm]) p3: 10.0;                  - with units
        # ((Number [degrees]) support_angle: 30.0;  - with units
        # (Number) hole_count: 5.0;                 - without units (unitless)
        pattern = r'[#*\(]*\((\w+)(?:\s*\[([^\]]*)\])?\)\s*([a-zA-Z_][a-zA-Z0-9_]*)\s*:\s*([-+]?\d*\.?\d+(?:[eE][-+]?\d+)?)'

        for match in re.finditer(pattern, self.text_content):
            expr_type, units, name, value = match.groups()
            expressions.append({
                'name': name,
                'value': float(value),
                'units': units if units else '',  # Empty string if no units
                'type': expr_type
            })

        return expressions

    def export_expressions_to_exp(self, output_path: Optional[Path] = None) -> Path:
        """
        Export expressions to .exp file using NX journal.

        This is the RECOMMENDED method for reading expressions because it:
        - Captures ALL expressions (formulas, references, constants)
        - Includes unitless expressions
        - Uses NX's native export, ensuring 100% accuracy

        Args:
            output_path: Path for .exp file (without .exp extension)
                        If None, uses temp file in same directory as .prt

        Returns:
            Path to the .exp file created
        """
        if output_path is None:
            # Create temp file in same directory
            output_path = self.prt_path.with_suffix('')  # Remove .prt
            output_path = Path(str(output_path) + "_expressions")

        # Get paths
        journal_script = Path(__file__).parent / "export_expressions.py"

        if not journal_script.exists():
            raise FileNotFoundError(f"Export journal script not found: {journal_script}")

        if not self.nx_run_journal_path.exists():
            raise FileNotFoundError(f"NX run_journal.exe not found: {self.nx_run_journal_path}")

        # Run NX journal to export expressions
        print(f"[NX] Exporting expressions from {self.prt_path.name} to .exp format...")

        # NX run_journal.exe syntax: run_journal.exe <journal-file> -args <arg1> <arg2> ...
        # Build command string with proper quoting
        cmd_str = f'"{self.nx_run_journal_path}" "{journal_script}" -args "{self.prt_path}" "{output_path}"'

        result = subprocess.run(cmd_str, capture_output=True, text=True, shell=True)

        exp_file = Path(str(output_path) + ".exp")

        # NOTE: NX run_journal.exe treats sys.exit(0) as a "syntax error" even though
        # it's a successful exit. We check if the file was created instead of return code.
        if not exp_file.exists():
            print(f"[ERROR] NX journal failed to create .exp file:")
            print(result.stdout)
            print(result.stderr)
            raise FileNotFoundError(f"Expected .exp file not created: {exp_file}")

        print(f"[OK] Expressions exported to: {exp_file}")
        return exp_file

    def parse_exp_file(self, exp_file_path: Path) -> Dict[str, Dict[str, any]]:
        """
        Parse a .exp file and return all expressions.

        .exp format examples:
            [MilliMeter]beam_length=5000
            [Kilogram]p173=973.968443678471
            hole_count=10
            Pattern_p7=hole_count

        Args:
            exp_file_path: Path to .exp file

        Returns:
            Dict mapping expression name to info dict with 'value', 'units', 'formula'
        """
        expressions = {}

        with open(exp_file_path, 'r', encoding='utf-8') as f:
            for line in f:
                line = line.strip()

                # Skip empty lines and comments
                if not line or line.startswith('//'):
                    continue

                # Pattern: [Unit]name=value or name=value
                # [MilliMeter]beam_length=5000
                # hole_count=10
                # Pattern_p7=hole_count (formula reference)
                match = re.match(r'(?:\[([^\]]+)\])?([a-zA-Z_][a-zA-Z0-9_]*)=(.*)', line)

                if match:
                    units, name, value_str = match.groups()

                    # Try to parse as number
                    try:
                        value = float(value_str)
                        formula = None
                    except ValueError:
                        # It's a formula/reference (e.g., "hole_count")
                        value = None
                        formula = value_str

                    expressions[name] = {
                        'value': value,
                        'units': units if units else '',
                        'formula': formula,
                        'type': 'Number'  # All .exp expressions are Number type
                    }

        return expressions

    def get_all_expressions(self, use_exp_export: bool = True) -> Dict[str, Dict[str, any]]:
        """
        Get all expressions as a dictionary.

        Args:
            use_exp_export: If True, uses NX .exp export (RECOMMENDED)
                           If False, uses binary .prt parsing (may miss expressions)

        Returns:
            Dict mapping expression name to info dict with 'value', 'units', 'type', 'formula'
        """
        if use_exp_export:
            # Use NX native .exp export (captures ALL expressions)
            try:
                exp_file = self.export_expressions_to_exp()
                expressions = self.parse_exp_file(exp_file)

                # Clean up temp file
                exp_file.unlink()

                return expressions
            except Exception as e:
                print(f"[WARNING] .exp export failed: {e}")
                print("[WARNING] Falling back to binary .prt parsing...")
                # Fall through to binary parsing

        # Fallback: Binary .prt parsing
        expressions_list = self.find_expressions()
        return {
            expr['name']: {
                'value': expr['value'],
                'units': expr['units'],
                'type': expr['type'],
                'formula': None  # Binary .prt files don't have formulas accessible
            }
            for expr in expressions_list
        }

    def update_expression(self, name: str, new_value: float) -> bool:
        """
        Update a single expression value.

        Args:
            name: Expression name
            new_value: New value

        Returns:
            True if updated, False if not found
        """
        # Find the expression pattern
        # Match: (Type [units]) name: old_value; OR (Type) name: old_value; (unitless)
        # We need to be careful to match the exact name and preserve formatting

        # Pattern that captures the full expression line
        # Units are optional (unitless expressions like hole_count don't have [units])
        pattern = rf'([#*\(]*\(\w+(?:\s*\[[^\]]*\])?\)\s*)({re.escape(name)})\s*:\s*([-+]?\d*\.?\d+(?:[eE][-+]?\d+)?)'

        matches = list(re.finditer(pattern, self.text_content))

        if not matches:
            print(f"Warning: Expression '{name}' not found in .prt file")
            return False

        if len(matches) > 1:
            print(f"Warning: Multiple matches for '{name}', updating first occurrence")

        # Get the first match
        match = matches[0]
        prefix, expr_name, old_value = match.groups()

        # Format new value (preserve decimal places if possible)
        # Check if old value had decimal point
        if '.' in old_value or 'e' in old_value.lower():
            # Use same precision as old value
            decimal_places = len(old_value.split('.')[-1]) if '.' in old_value else 2
            new_value_str = f"{new_value:.{decimal_places}f}"
        else:
            # Integer format
            new_value_str = f"{int(new_value)}"

        # Build replacement string
        replacement = f"{prefix}{expr_name}: {new_value_str}"

        # Replace in text content
        old_match = match.group(0)
        self.text_content = self.text_content.replace(old_match, replacement, 1)

        # Also update in binary content
        old_bytes = old_match.encode('latin-1')
        new_bytes = replacement.encode('latin-1')

        # Find and replace in binary content
        start_pos = self.content.find(old_bytes)
        if start_pos != -1:
            # Replace bytes
            self.content[start_pos:start_pos+len(old_bytes)] = new_bytes

            print(f"Updated: {name} = {old_value} -> {new_value_str}")
            return True
        else:
            print(f"Warning: Could not update binary content for '{name}'")
            return False

    def update_expressions(self, updates: Dict[str, float], use_nx_import: bool = True):
        """
        Update multiple expressions at once.

        Args:
            updates: Dict mapping expression name to new value
                    {'tip_thickness': 22.5, 'support_angle': 35.0}
            use_nx_import: If True, uses NX journal to import .exp file (RECOMMENDED for all expressions)
                          If False, uses binary .prt editing (may miss some expressions)
        """
        if use_nx_import:
            # Use NX journal to import expressions
            return self.update_expressions_via_import(updates)

        # Fallback: Binary .prt editing
        print(f"\nUpdating {len(updates)} expressions in {self.prt_path.name}:")

        updated_count = 0
        for name, value in updates.items():
            if self.update_expression(name, value):
                updated_count += 1

        print(f"Successfully updated {updated_count}/{len(updates)} expressions")

    def update_expressions_via_import(self, updates: Dict[str, float]):
        """
        Update expressions by creating a .exp file and importing it via NX journal.

        This method works for ALL expressions including those not stored in text format
        in the binary .prt file (like hole_count).

        Args:
            updates: Dict mapping expression name to new value
        """
        print(f"\nUpdating {len(updates)} expressions via NX .exp import:")

        # Get all expressions to determine units
        all_expressions = self.get_all_expressions(use_exp_export=True)

        # Create .exp file with ONLY the study variables
        exp_file = self.prt_path.parent / f"{self.prt_path.stem}_study_variables.exp"

        with open(exp_file, 'w', encoding='utf-8') as f:
            for name, value in updates.items():
                if name in all_expressions:
                    units = all_expressions[name].get('units', '')
                    if units:
                        # Expression with units: [MilliMeter]beam_length=5000
                        f.write(f"[{units}]{name}={value}\n")
                    else:
                        # Unitless expression: hole_count=10
                        f.write(f"{name}={value}\n")
                    print(f"  {name}: {value} {units if units else ''}")
                else:
                    print(f"  Warning: {name} not found in part expressions, skipping")

        print(f"\n[EXP] Created: {exp_file}")

        # Run NX journal to import expressions
        journal_script = Path(__file__).parent / "import_expressions.py"

        if not journal_script.exists():
            raise FileNotFoundError(f"Import journal script not found: {journal_script}")

        if not self.nx_run_journal_path.exists():
            raise FileNotFoundError(f"NX run_journal.exe not found: {self.nx_run_journal_path}")

        print(f"[NX] Importing expressions into {self.prt_path.name}...")

        # Build command
        cmd_str = f'"{self.nx_run_journal_path}" "{journal_script}" -args "{self.prt_path}" "{exp_file}"'

        result = subprocess.run(cmd_str, capture_output=True, text=True, shell=True)

        # Clean up .exp file
        exp_file.unlink()

        # Check if import succeeded
        if result.returncode != 0 and "successfully" not in result.stdout.lower():
            print(f"[ERROR] NX journal failed:")
            print(result.stdout)
            print(result.stderr)
            raise RuntimeError(f"Expression import failed")

        print(f"[OK] All {len(updates)} expressions updated successfully!")

    def save(self, output_path: Path = None):
        """
        Save modified .prt file.

        Args:
            output_path: Optional different path to save to.
                        If None, overwrites original (with backup if enabled)
        """
        if output_path is None:
            output_path = self.prt_path
            if self.backup_enabled:
                self._create_backup()

        # Write updated binary content
        with open(output_path, 'wb') as f:
            f.write(self.content)

        print(f"Saved to: {output_path}")

    def verify_update(self, name: str, expected_value: float, tolerance: float = 1e-6) -> bool:
        """
        Verify that an expression was updated correctly.

        Args:
            name: Expression name
            expected_value: Expected value
            tolerance: Acceptable difference

        Returns:
            True if value matches (within tolerance)
        """
        expressions = self.find_expressions()
        expr = next((e for e in expressions if e['name'] == name), None)

        if expr is None:
            print(f"Expression '{name}' not found")
            return False

        actual_value = expr['value']
        difference = abs(actual_value - expected_value)

        if difference <= tolerance:
            print(f"OK Verified: {name} = {actual_value} (expected {expected_value})")
            return True
        else:
            print(f"FAIL Verification failed: {name} = {actual_value}, expected {expected_value} (diff: {difference})")
            return False


# Convenience function for optimization loop
def update_nx_model(prt_file_path: Path, design_variables: Dict[str, float], backup: bool = False):
    """
    Convenience function to update NX model parameters.

    Args:
        prt_file_path: Path to .prt file
        design_variables: Dict of parameter name -> value
        backup: Whether to create backup

    Example:
        >>> update_nx_model(
        ...     Path("Bracket.prt"),
        ...     {'tip_thickness': 22.5, 'support_angle': 35.0}
        ... )
    """
    updater = NXParameterUpdater(prt_file_path, backup=backup)
    updater.update_expressions(design_variables)
    updater.save()


# Example usage
if __name__ == "__main__":
    import sys

    if len(sys.argv) < 2:
        print("Usage: python nx_updater.py <path_to_prt_file>")
        sys.exit(1)

    prt_path = Path(sys.argv[1])

    # Test: Find and print all expressions
    print("="*60)
    print("NX PARAMETER UPDATER TEST")
    print("="*60)

    updater = NXParameterUpdater(prt_path, backup=True)

    print("\nCurrent expressions in file:")
    expressions = updater.find_expressions()
    for expr in expressions:
        print(f"  {expr['name']}: {expr['value']} {expr['units']}")

    # Test update (if expressions found)
    if expressions:
        print("\n" + "="*60)
        print("TEST UPDATE")
        print("="*60)

        # Update first expression
        first_expr = expressions[0]
        test_name = first_expr['name']
        test_new_value = first_expr['value'] * 1.1  # Increase by 10%

        print(f"\nUpdating {test_name} from {first_expr['value']} to {test_new_value}")

        updater.update_expression(test_name, test_new_value)

        # Save to test file
        test_output = prt_path.with_suffix('.prt.test')
        updater.save(test_output)

        # Verify by re-reading
        print("\n" + "="*60)
        print("VERIFICATION")
        print("="*60)
        verifier = NXParameterUpdater(test_output, backup=False)
        verifier.verify_update(test_name, test_new_value)

        print(f"\nTest complete. Modified file: {test_output}")
    else:
        print("\nNo expressions found in file. Nothing to test.")