Atomizer/optimization_engine/validators/model_validator.py

"""
Model Validator for Atomizer
============================

Validates NX model files and simulation setup before running optimizations.
Checks file existence, structure, and configuration compatibility.

Usage:
    from optimization_engine.validators import validate_model, validate_model_files

    # Validate model directory
    result = validate_model("studies/my_study/1_setup/model")

    # Validate specific files
    result = validate_model_files(
        prt_file="Beam.prt",
        sim_file="Beam_sim1.sim",
        model_dir=Path("studies/my_study/1_setup/model")
    )

    if result.is_valid:
        print("Model is ready!")
    else:
        for error in result.errors:
            print(f"ERROR: {error}")
"""

import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import List, Dict, Any, Optional, Union


@dataclass
class ModelError:
    """Represents a model validation error that blocks execution."""
    component: str
    message: str
    suggestion: Optional[str] = None

    def __str__(self):
        msg = f"[{self.component}] {self.message}"
        if self.suggestion:
            msg += f" (Suggestion: {self.suggestion})"
        return msg


@dataclass
class ModelWarning:
    """Represents a model validation warning."""
    component: str
    message: str
    suggestion: Optional[str] = None

    def __str__(self):
        msg = f"[{self.component}] {self.message}"
        if self.suggestion:
            msg += f" (Suggestion: {self.suggestion})"
        return msg


@dataclass
class ModelValidationResult:
    """Result of model validation."""
    errors: List[ModelError] = field(default_factory=list)
    warnings: List[ModelWarning] = field(default_factory=list)

    # Discovered files
    prt_file: Optional[Path] = None
    sim_file: Optional[Path] = None
    fem_file: Optional[Path] = None

    # Model info
    model_name: Optional[str] = None
    model_dir: Optional[Path] = None
    file_sizes: Dict[str, int] = field(default_factory=dict)

    @property
    def is_valid(self) -> bool:
        """Model is valid if there are no errors."""
        return len(self.errors) == 0

    @property
    def has_simulation(self) -> bool:
        """Check if simulation file exists."""
        return self.sim_file is not None

    @property
    def has_fem(self) -> bool:
        """Check if FEM mesh file exists."""
        return self.fem_file is not None

    def __str__(self):
        lines = []
        lines.append(f"Model: {self.model_name or 'Unknown'}")
        lines.append(f"Directory: {self.model_dir or 'Unknown'}")
        lines.append("")

        lines.append("Files:")
        if self.prt_file:
            size = self.file_sizes.get('prt', 0)
            lines.append(f"  [OK] Part file: {self.prt_file.name} ({_format_size(size)})")
        else:
            lines.append("  [X] Part file: NOT FOUND")

        if self.sim_file:
            size = self.file_sizes.get('sim', 0)
            lines.append(f"  [OK] Simulation: {self.sim_file.name} ({_format_size(size)})")
        else:
            lines.append("  [X] Simulation: NOT FOUND")

        if self.fem_file:
            size = self.file_sizes.get('fem', 0)
            lines.append(f"  [OK] FEM mesh: {self.fem_file.name} ({_format_size(size)})")
        else:
            lines.append("  ? FEM mesh: Not found (will be created on first solve)")

        if self.errors:
            lines.append("")
            lines.append(f"ERRORS ({len(self.errors)}):")
            for e in self.errors:
                lines.append(f"  - {e}")

        if self.warnings:
            lines.append("")
            lines.append(f"WARNINGS ({len(self.warnings)}):")
            for w in self.warnings:
                lines.append(f"  - {w}")

        if self.is_valid:
            lines.append("")
            lines.append("[OK] Model validation passed!")

        return "\n".join(lines)


def _format_size(size_bytes: int) -> str:
    """Format file size for display."""
    if size_bytes < 1024:
        return f"{size_bytes} B"
    elif size_bytes < 1024 * 1024:
        return f"{size_bytes / 1024:.1f} KB"
    else:
        return f"{size_bytes / (1024 * 1024):.1f} MB"


def validate_model(model_dir: Union[str, Path],
                  expected_model_name: Optional[str] = None) -> ModelValidationResult:
    """
    Validate an NX model directory.

    Args:
        model_dir: Path to the model directory
        expected_model_name: Expected base name of the model (optional)

    Returns:
        ModelValidationResult with errors, warnings, and discovered files
    """
    model_dir = Path(model_dir)
    result = ModelValidationResult(model_dir=model_dir)

    # Check directory exists
    if not model_dir.exists():
        result.errors.append(ModelError(
            component="directory",
            message=f"Model directory not found: {model_dir}",
            suggestion="Create the directory and add NX model files"
        ))
        return result

    if not model_dir.is_dir():
        result.errors.append(ModelError(
            component="directory",
            message=f"Path is not a directory: {model_dir}",
            suggestion="Provide path to the model directory, not a file"
        ))
        return result

    # Find model files
    prt_files = list(model_dir.glob("*.prt"))
    sim_files = list(model_dir.glob("*.sim"))
    fem_files = list(model_dir.glob("*.fem"))

    # Check for part file
    if len(prt_files) == 0:
        result.errors.append(ModelError(
            component="part",
            message="No .prt file found in model directory",
            suggestion="Add your NX part file to the model directory"
        ))
    elif len(prt_files) > 1:
        # Filter out internal files (often have _i suffix)
        main_prt_files = [f for f in prt_files if not f.stem.endswith('_i')]
        if len(main_prt_files) == 1:
            prt_files = main_prt_files
        elif expected_model_name:
            matching = [f for f in prt_files if f.stem == expected_model_name]
            if matching:
                prt_files = matching
            else:
                result.warnings.append(ModelWarning(
                    component="part",
                    message=f"Multiple .prt files found, none match expected name '{expected_model_name}'",
                    suggestion="Specify the correct model name in configuration"
                ))
        else:
            result.warnings.append(ModelWarning(
                component="part",
                message=f"Multiple .prt files found: {[f.name for f in prt_files]}",
                suggestion="Consider keeping only the main model file in the directory"
            ))

    if prt_files:
        result.prt_file = prt_files[0]
        result.model_name = result.prt_file.stem
        result.file_sizes['prt'] = result.prt_file.stat().st_size

        # Validate part file
        _validate_prt_file(result.prt_file, result)

    # Check for simulation file
    if len(sim_files) == 0:
        result.errors.append(ModelError(
            component="simulation",
            message="No .sim file found in model directory",
            suggestion="Create a simulation in NX and save it to this directory"
        ))
    elif len(sim_files) > 1:
        if result.model_name:
            # Try to find matching sim file
            expected_sim = f"{result.model_name}_sim1.sim"
            matching = [f for f in sim_files if f.name.lower() == expected_sim.lower()]
            if matching:
                sim_files = matching
            else:
                result.warnings.append(ModelWarning(
                    component="simulation",
                    message=f"Multiple .sim files found: {[f.name for f in sim_files]}",
                    suggestion=f"Expected: {expected_sim}"
                ))
        else:
            result.warnings.append(ModelWarning(
                component="simulation",
                message=f"Multiple .sim files found: {[f.name for f in sim_files]}",
                suggestion="Keep only one simulation file"
            ))

    if sim_files:
        result.sim_file = sim_files[0]
        result.file_sizes['sim'] = result.sim_file.stat().st_size

        # Validate simulation file
        _validate_sim_file(result.sim_file, result)

    # Check for FEM file
    if len(fem_files) == 0:
        result.warnings.append(ModelWarning(
            component="fem",
            message="No .fem file found",
            suggestion="FEM mesh will be created automatically on first solve"
        ))
    else:
        if result.model_name:
            expected_fem = f"{result.model_name}_fem1.fem"
            matching = [f for f in fem_files if f.name.lower() == expected_fem.lower()]
            if matching:
                fem_files = matching

        result.fem_file = fem_files[0]
        result.file_sizes['fem'] = result.fem_file.stat().st_size

    # Cross-validate files
    _validate_file_relationships(result)

    return result


def validate_model_files(prt_file: Union[str, Path],
                        sim_file: Union[str, Path],
                        model_dir: Optional[Union[str, Path]] = None) -> ModelValidationResult:
    """
    Validate specific model files.

    Args:
        prt_file: Name or path to the part file
        sim_file: Name or path to the simulation file
        model_dir: Base directory (optional, will be inferred if full paths given)

    Returns:
        ModelValidationResult
    """
    prt_path = Path(prt_file)
    sim_path = Path(sim_file)

    # If paths are relative and model_dir provided, resolve them
    if model_dir:
        model_dir = Path(model_dir)
        if not prt_path.is_absolute():
            prt_path = model_dir / prt_path
        if not sim_path.is_absolute():
            sim_path = model_dir / sim_path
    else:
        # Infer model_dir from prt_file
        if prt_path.is_absolute():
            model_dir = prt_path.parent
        else:
            model_dir = Path.cwd()

    result = ModelValidationResult(model_dir=model_dir)

    # Check part file
    if not prt_path.exists():
        result.errors.append(ModelError(
            component="part",
            message=f"Part file not found: {prt_path}",
            suggestion="Check the file path and name"
        ))
    else:
        result.prt_file = prt_path
        result.model_name = prt_path.stem
        result.file_sizes['prt'] = prt_path.stat().st_size
        _validate_prt_file(prt_path, result)

    # Check simulation file
    if not sim_path.exists():
        result.errors.append(ModelError(
            component="simulation",
            message=f"Simulation file not found: {sim_path}",
            suggestion="Check the file path and name"
        ))
    else:
        result.sim_file = sim_path
        result.file_sizes['sim'] = sim_path.stat().st_size
        _validate_sim_file(sim_path, result)

    # Check for FEM file
    if result.model_name:
        fem_path = model_dir / f"{result.model_name}_fem1.fem"
        if fem_path.exists():
            result.fem_file = fem_path
            result.file_sizes['fem'] = fem_path.stat().st_size
        else:
            # Try alternative naming
            fem_files = list(model_dir.glob("*.fem")) if model_dir.exists() else []
            if fem_files:
                result.fem_file = fem_files[0]
                result.file_sizes['fem'] = result.fem_file.stat().st_size

    _validate_file_relationships(result)

    return result


def _validate_prt_file(prt_path: Path, result: ModelValidationResult):
    """Validate a part file."""
    # Check file size
    size = prt_path.stat().st_size

    if size == 0:
        result.errors.append(ModelError(
            component="part",
            message="Part file is empty",
            suggestion="Re-save the part file in NX"
        ))
        return

    if size < 1024:
        result.warnings.append(ModelWarning(
            component="part",
            message=f"Part file is very small ({_format_size(size)})",
            suggestion="Verify the file contains valid geometry"
        ))

    # Check for NX file signature (basic validation)
    try:
        with open(prt_path, 'rb') as f:
            header = f.read(8)
            # NX files typically start with a specific signature
            # This is a basic check - real NX files have more complex headers
            if len(header) < 8:
                result.warnings.append(ModelWarning(
                    component="part",
                    message="Part file appears incomplete",
                    suggestion="Re-save the file in NX"
                ))
    except PermissionError:
        result.errors.append(ModelError(
            component="part",
            message="Cannot read part file - permission denied",
            suggestion="Close NX if the file is open, or check file permissions"
        ))
    except Exception as e:
        result.warnings.append(ModelWarning(
            component="part",
            message=f"Could not verify part file: {e}",
            suggestion="Ensure file is a valid NX part"
        ))


def _validate_sim_file(sim_path: Path, result: ModelValidationResult):
    """Validate a simulation file."""
    size = sim_path.stat().st_size

    if size == 0:
        result.errors.append(ModelError(
            component="simulation",
            message="Simulation file is empty",
            suggestion="Re-save the simulation in NX"
        ))
        return

    if size < 512:
        result.warnings.append(ModelWarning(
            component="simulation",
            message=f"Simulation file is very small ({_format_size(size)})",
            suggestion="Verify simulation setup in NX"
        ))

    # Check for assembly FEM - detect missing assembly .prt file
    _validate_assembly_components(sim_path, result)


def _validate_assembly_components(sim_path: Path, result: ModelValidationResult):
    """
    Validate that all .prt files required by an assembly FEM are present.

    Assembly FEM (.afm) files reference:
    - The assembly .prt file (e.g., ASSY_M1.prt)
    - Component part files (e.g., M1_Blank.prt)
    - Component FEM idealized parts (e.g., M1_Blank_fem1_i.prt)

    This function detects when the main assembly .prt file is missing.
    """
    model_dir = sim_path.parent
    sim_stem = sim_path.stem  # e.g., "ASSY_M1_assyfem1_sim1"

    # Check if this is an assembly FEM by looking for _assyfem pattern
    if '_assyfem' not in sim_stem.lower():
        return  # Not an assembly FEM

    # Extract assembly name from sim file: ASSY_M1_assyfem1_sim1 -> ASSY_M1
    # Pattern: {ASSY_NAME}_assyfem{N}_sim{M}.sim
    import re
    match = re.match(r'^(.+?)_assyfem\d+', sim_stem, re.IGNORECASE)
    if not match:
        return

    assembly_name = match.group(1)  # e.g., "ASSY_M1"
    expected_assy_prt = f"{assembly_name}.prt"
    expected_assy_path = model_dir / expected_assy_prt

    # Check if assembly .prt file exists
    if not expected_assy_path.exists():
        # Try case-insensitive search
        all_prts = list(model_dir.glob("*.prt"))
        matching = [p for p in all_prts if p.stem.lower() == assembly_name.lower()]

        if not matching:
            result.errors.append(ModelError(
                component="assembly",
                message=f"CRITICAL: Assembly part file missing: {expected_assy_prt}",
                suggestion=f"Copy {expected_assy_prt} from the source model directory to {model_dir}"
            ))

            # Also warn about potential related files
            result.warnings.append(ModelWarning(
                component="assembly",
                message=f"Assembly FEM requires all component .prt and .fem files",
                suggestion="Ensure ALL .prt, .fem, and _i.prt files from the source are copied"
            ))
    else:
        # Assembly .prt exists - verify it's readable
        if not os.access(expected_assy_path, os.R_OK):
            result.errors.append(ModelError(
                component="assembly",
                message=f"Assembly part file not readable: {expected_assy_prt}",
                suggestion="Check file permissions"
            ))

    # Also check for the .afm file
    afm_name = f"{assembly_name}_assyfem1.afm"  # Common pattern
    # Try to find any .afm file
    afm_files = list(model_dir.glob("*.afm"))
    if not afm_files:
        result.warnings.append(ModelWarning(
            component="assembly",
            message="No .afm file found for assembly FEM",
            suggestion="The .afm file may be created on first solve"
        ))


def _validate_file_relationships(result: ModelValidationResult):
    """Validate relationships between model files."""
    if not result.prt_file or not result.sim_file:
        return

    # Check naming convention
    prt_stem = result.prt_file.stem
    sim_stem = result.sim_file.stem

    expected_sim_stem = f"{prt_stem}_sim1"
    if sim_stem != expected_sim_stem and not sim_stem.startswith(prt_stem):
        result.warnings.append(ModelWarning(
            component="naming",
            message=f"Simulation name '{sim_stem}' doesn't match part name '{prt_stem}'",
            suggestion=f"Expected simulation name: {expected_sim_stem}.sim"
        ))

    # Check FEM naming if present
    if result.fem_file:
        fem_stem = result.fem_file.stem
        expected_fem_stem = f"{prt_stem}_fem1"
        if fem_stem != expected_fem_stem and not fem_stem.startswith(prt_stem):
            result.warnings.append(ModelWarning(
                component="naming",
                message=f"FEM name '{fem_stem}' doesn't match part name '{prt_stem}'",
                suggestion=f"Expected FEM name: {expected_fem_stem}.fem"
            ))

    # Check files are in same directory
    if result.prt_file.parent != result.sim_file.parent:
        result.warnings.append(ModelWarning(
            component="directory",
            message="Part and simulation files are in different directories",
            suggestion="Keep all model files in the same directory"
        ))


def validate_study_model(study_name: str,
                        studies_dir: str = "studies",
                        config: Optional[Dict[str, Any]] = None) -> ModelValidationResult:
    """
    Validate model for a complete study.

    Args:
        study_name: Name of the study folder (e.g., "uav_arm_optimization")
        studies_dir: Base directory for studies (default: "studies")
        config: Optional optimization_config.json contents (loaded dict, not path)

    Returns:
        ModelValidationResult
    """
    study_path = Path(studies_dir) / study_name
    model_dir = study_path / "1_setup" / "model"

    # Load config if not provided
    if config is None:
        config_path = study_path / "1_setup" / "optimization_config.json"
        if config_path.exists():
            import json
            try:
                with open(config_path, 'r') as f:
                    config = json.load(f)
            except (json.JSONDecodeError, IOError):
                config = None

    # Get expected file names from config if available
    expected_model_name = None
    if config and isinstance(config, dict) and 'simulation' in config:
        sim_config = config['simulation']
        if 'model_file' in sim_config:
            expected_model_name = Path(sim_config['model_file']).stem

    result = validate_model(model_dir, expected_model_name)

    # Additional study-specific validations
    if config and isinstance(config, dict):
        _validate_config_model_match(config, result)

    return result


def _validate_config_model_match(config: Dict[str, Any], result: ModelValidationResult):
    """Check that config matches discovered model files."""
    sim_config = config.get('simulation', {})

    # Check model file name matches
    if 'model_file' in sim_config and result.prt_file:
        config_model = Path(sim_config['model_file']).name
        actual_model = result.prt_file.name
        if config_model.lower() != actual_model.lower():
            result.warnings.append(ModelWarning(
                component="config",
                message=f"Config specifies '{config_model}' but found '{actual_model}'",
                suggestion="Update config to match actual file name"
            ))

    # Check sim file name matches
    if 'sim_file' in sim_config and result.sim_file:
        config_sim = Path(sim_config['sim_file']).name
        actual_sim = result.sim_file.name
        if config_sim.lower() != actual_sim.lower():
            result.warnings.append(ModelWarning(
                component="config",
                message=f"Config specifies '{config_sim}' but found '{actual_sim}'",
                suggestion="Update config to match actual file name"
            ))


# CLI interface for direct execution
if __name__ == "__main__":
    import sys

    if len(sys.argv) < 2:
        print("Usage: python model_validator.py <path_to_model_directory>")
        print("       python model_validator.py <path_to_study_directory>")
        sys.exit(1)

    path = Path(sys.argv[1])

    # Check if it's a study directory or model directory
    if (path / "1_setup" / "model").exists():
        # It's a study directory
        result = validate_study_model(path)
    elif path.is_dir():
        # It's a model directory
        result = validate_model(path)
    else:
        print(f"ERROR: Path not found or not a directory: {path}")
        sys.exit(1)

    print(result)

    if result.is_valid:
        print("\n✓ Model validation passed!")
        sys.exit(0)
    else:
        print(f"\n✗ Model has {len(result.errors)} error(s)")
        sys.exit(1)