Atomizer/optimization_engine/nx_updater.py

"""
NX Parameter Updater

Updates design variable values in NX .prt files.

NX .prt files are binary, but expressions are stored in readable text sections.
This module can update expression values by:
1. Reading the binary file
2. Finding and replacing expression value patterns
3. Writing back the updated file

Alternative: Use NXOpen API if NX is running (future enhancement)
"""

from pathlib import Path
from typing import Dict, List
import re
import shutil
from datetime import datetime


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

    NX Expression Format in binary .prt files:
    #(Number [mm]) tip_thickness: 20.0;
    *(Number [degrees]) support_angle: 30.0;
    """

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

        Args:
            prt_file_path: Path to NX .prt file
            backup: If True, create backup before modifying
        """
        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
        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:
        # #(Number [mm]) tip_thickness: 20.0;
        # *(Number [mm]) p3: 10.0;
        # ((Number [degrees]) support_angle: 30.0;
        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,
                'type': expr_type
            })

        return expressions

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

        Returns:
            Dict mapping expression name to info dict with 'value', 'units', 'type'
        """
        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;
        # We need to be careful to match the exact name and preserve formatting

        # Pattern that captures the full expression line
        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]):
        """
        Update multiple expressions at once.

        Args:
            updates: Dict mapping expression name to new value
                    {'tip_thickness': 22.5, 'support_angle': 35.0}
        """
        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 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.")