Atomizer/optimization_engine/utils/nx_file_discovery.py

"""
NX File Discovery Utility
=========================

Discovers all files required by an NX simulation (.sim) file and provides
functions to copy them to a target directory.

For assembly FEM models, this includes:
- The .sim file itself
- The .afm (Assembly FEM) file
- The assembly .prt file
- All component .prt files
- All idealized part files (*_i.prt)
- All component .fem files

Usage:
    from optimization_engine.utils.nx_file_discovery import (
        discover_sim_dependencies,
        copy_sim_with_dependencies
    )

    # Discover all required files
    deps = discover_sim_dependencies("/path/to/source/model.sim")
    print(f"Found {len(deps.all_files)} required files")

    # Copy all files to target directory
    result = copy_sim_with_dependencies(
        sim_file="/path/to/source/model.sim",
        target_dir="/path/to/target/model/"
    )
"""

import os
import re
import shutil
from dataclasses import dataclass, field
from pathlib import Path
from typing import List, Set, Optional, Dict


@dataclass
class SimDependencies:
    """Represents all files required by a .sim file."""
    sim_file: Path
    source_dir: Path

    # Core files
    afm_file: Optional[Path] = None
    assembly_prt: Optional[Path] = None

    # Component files
    component_prts: List[Path] = field(default_factory=list)
    idealized_prts: List[Path] = field(default_factory=list)
    fem_files: List[Path] = field(default_factory=list)

    # Additional output files (optional, not always needed)
    output_files: List[Path] = field(default_factory=list)

    @property
    def all_files(self) -> List[Path]:
        """Get list of all required files."""
        files = [self.sim_file]
        if self.afm_file:
            files.append(self.afm_file)
        if self.assembly_prt:
            files.append(self.assembly_prt)
        files.extend(self.component_prts)
        files.extend(self.idealized_prts)
        files.extend(self.fem_files)
        return files

    @property
    def all_files_with_outputs(self) -> List[Path]:
        """Get all files including output files."""
        return self.all_files + self.output_files

    def __str__(self):
        lines = [
            f"Simulation Dependencies for: {self.sim_file.name}",
            f"Source directory: {self.source_dir}",
            "",
            "Required files:"
        ]

        if self.afm_file:
            lines.append(f"  [AFM] {self.afm_file.name}")
        if self.assembly_prt:
            lines.append(f"  [ASSY] {self.assembly_prt.name}")
        for prt in self.component_prts:
            lines.append(f"  [PRT] {prt.name}")
        for prt in self.idealized_prts:
            lines.append(f"  [PRT_i] {prt.name}")
        for fem in self.fem_files:
            lines.append(f"  [FEM] {fem.name}")

        lines.append(f"\nTotal: {len(self.all_files)} required files")
        return "\n".join(lines)


def discover_sim_dependencies(sim_file: Path,
                              include_outputs: bool = False) -> SimDependencies:
    """
    Discover all files required by an NX simulation file.

    Args:
        sim_file: Path to the .sim file
        include_outputs: Whether to include output files (.op2, .f06, etc.)

    Returns:
        SimDependencies object with all discovered files
    """
    sim_file = Path(sim_file)
    if not sim_file.exists():
        raise FileNotFoundError(f"Simulation file not found: {sim_file}")

    source_dir = sim_file.parent
    sim_stem = sim_file.stem  # e.g., "ASSY_M1_assyfem1_sim1"

    deps = SimDependencies(sim_file=sim_file, source_dir=source_dir)

    # Detect if this is an assembly FEM
    is_assembly = '_assyfem' in sim_stem.lower()

    if is_assembly:
        _discover_assembly_deps(deps, sim_stem, source_dir)
    else:
        _discover_single_part_deps(deps, sim_stem, source_dir)

    if include_outputs:
        _discover_output_files(deps, sim_stem, source_dir)

    return deps


def _discover_assembly_deps(deps: SimDependencies, sim_stem: str, source_dir: Path):
    """Discover dependencies for an assembly FEM."""
    # Extract assembly name: ASSY_M1_assyfem1_sim1 -> ASSY_M1
    match = re.match(r'^(.+?)_assyfem(\d+)', sim_stem, re.IGNORECASE)
    if not match:
        return

    assembly_name = match.group(1)  # e.g., "ASSY_M1"
    afm_num = match.group(2)  # e.g., "1"

    # Look for the .afm file
    afm_pattern = f"{assembly_name}_assyfem{afm_num}.afm"
    afm_files = list(source_dir.glob(afm_pattern)) or \
                [f for f in source_dir.glob("*.afm") if assembly_name.lower() in f.stem.lower()]
    if afm_files:
        deps.afm_file = afm_files[0]

    # Look for the assembly .prt file
    assy_prt = source_dir / f"{assembly_name}.prt"
    if assy_prt.exists():
        deps.assembly_prt = assy_prt
    else:
        # Try case-insensitive search
        for prt in source_dir.glob("*.prt"):
            if prt.stem.lower() == assembly_name.lower():
                deps.assembly_prt = prt
                break

    # Discover component files based on existing .fem files in the directory
    # Pattern: Component FEM files are typically {ComponentName}_fem{N}.fem
    for fem_file in source_dir.glob("*.fem"):
        deps.fem_files.append(fem_file)

        # Extract component name from FEM: M1_Blank_fem1.fem -> M1_Blank
        fem_match = re.match(r'^(.+?)_fem\d+$', fem_file.stem, re.IGNORECASE)
        if fem_match:
            component_name = fem_match.group(1)

            # Look for component .prt
            comp_prt = source_dir / f"{component_name}.prt"
            if comp_prt.exists() and comp_prt not in deps.component_prts:
                deps.component_prts.append(comp_prt)

            # Look for idealized part (*_fem{N}_i.prt)
            idealized_prt = source_dir / f"{fem_file.stem}_i.prt"
            if idealized_prt.exists() and idealized_prt not in deps.idealized_prts:
                deps.idealized_prts.append(idealized_prt)

    # Also scan for any _i.prt files that might have been missed
    for prt in source_dir.glob("*_i.prt"):
        if prt not in deps.idealized_prts:
            deps.idealized_prts.append(prt)


def _discover_single_part_deps(deps: SimDependencies, sim_stem: str, source_dir: Path):
    """Discover dependencies for a single-part simulation."""
    # Pattern: Model_sim1.sim -> Model.prt, Model_fem1.fem
    match = re.match(r'^(.+?)_sim\d+$', sim_stem, re.IGNORECASE)
    if not match:
        # Try without _sim suffix
        base_name = sim_stem
    else:
        base_name = match.group(1)

    # Look for main .prt file
    main_prt = source_dir / f"{base_name}.prt"
    if main_prt.exists():
        deps.component_prts.append(main_prt)
    else:
        # Try to find any matching prt
        for prt in source_dir.glob("*.prt"):
            if base_name.lower() in prt.stem.lower():
                deps.component_prts.append(prt)

    # Look for .fem file
    for fem in source_dir.glob(f"{base_name}_fem*.fem"):
        deps.fem_files.append(fem)

    # Look for idealized parts
    for prt in source_dir.glob(f"{base_name}_fem*_i.prt"):
        deps.idealized_prts.append(prt)


def _discover_output_files(deps: SimDependencies, sim_stem: str, source_dir: Path):
    """Discover output files (.op2, .f06, etc.)."""
    # Common output patterns
    output_extensions = ['.op2', '.f06', '.f04', '.log', '.dat', '.diag', '.csv']

    for ext in output_extensions:
        for f in source_dir.glob(f"*{ext}"):
            if sim_stem.lower() in f.stem.lower() or \
               deps.sim_file.stem.lower().replace('_sim', '') in f.stem.lower():
                deps.output_files.append(f)


@dataclass
class CopyResult:
    """Result of copying sim dependencies."""
    success: bool
    copied_files: List[Path] = field(default_factory=list)
    failed_files: List[tuple] = field(default_factory=list)  # (Path, error message)
    target_dir: Optional[Path] = None

    def __str__(self):
        lines = [f"Copy Result: {'SUCCESS' if self.success else 'FAILED'}"]
        if self.target_dir:
            lines.append(f"Target: {self.target_dir}")
        lines.append(f"Copied: {len(self.copied_files)} files")
        if self.failed_files:
            lines.append(f"Failed: {len(self.failed_files)} files")
            for path, err in self.failed_files:
                lines.append(f"  - {path.name}: {err}")
        return "\n".join(lines)


def copy_sim_with_dependencies(sim_file: Path,
                               target_dir: Path,
                               include_outputs: bool = False,
                               overwrite: bool = True) -> CopyResult:
    """
    Copy a simulation file and all its dependencies to a target directory.

    Args:
        sim_file: Path to the .sim file
        target_dir: Destination directory
        include_outputs: Whether to include output files (.op2, etc.)
        overwrite: Whether to overwrite existing files

    Returns:
        CopyResult with status and details
    """
    sim_file = Path(sim_file)
    target_dir = Path(target_dir)

    # Discover dependencies
    deps = discover_sim_dependencies(sim_file, include_outputs=include_outputs)

    # Create target directory if needed
    target_dir.mkdir(parents=True, exist_ok=True)

    result = CopyResult(success=True, target_dir=target_dir)

    # Get files to copy
    files_to_copy = deps.all_files_with_outputs if include_outputs else deps.all_files

    for src_file in files_to_copy:
        dst_file = target_dir / src_file.name

        try:
            if dst_file.exists():
                if not overwrite:
                    result.copied_files.append(dst_file)  # Count as success
                    continue
                # Remove read-only flag if needed
                if not os.access(dst_file, os.W_OK):
                    os.chmod(dst_file, 0o666)

            shutil.copy2(src_file, dst_file)

            # Make sure the copy is writable
            if not os.access(dst_file, os.W_OK):
                os.chmod(dst_file, 0o666)

            result.copied_files.append(dst_file)

        except Exception as e:
            result.failed_files.append((src_file, str(e)))
            result.success = False

    return result


def get_required_files_from_sim(sim_file: Path) -> List[str]:
    """
    Get a simple list of required file names for a simulation.

    Useful for validation and reporting.

    Args:
        sim_file: Path to the .sim file

    Returns:
        List of required file names (not full paths)
    """
    deps = discover_sim_dependencies(sim_file)
    return [f.name for f in deps.all_files]


# CLI interface
if __name__ == "__main__":
    import sys

    if len(sys.argv) < 2:
        print("Usage: python nx_file_discovery.py <sim_file> [target_dir]")
        print("       python nx_file_discovery.py --discover <sim_file>")
        sys.exit(1)

    if sys.argv[1] == "--discover":
        if len(sys.argv) < 3:
            print("Usage: python nx_file_discovery.py --discover <sim_file>")
            sys.exit(1)
        sim_file = Path(sys.argv[2])
        deps = discover_sim_dependencies(sim_file, include_outputs=True)
        print(deps)
    else:
        sim_file = Path(sys.argv[1])
        if len(sys.argv) >= 3:
            target_dir = Path(sys.argv[2])
            result = copy_sim_with_dependencies(sim_file, target_dir)
            print(result)
        else:
            # Just discover
            deps = discover_sim_dependencies(sim_file)
            print(deps)