Atomizer/optimization_engine/result_extractors/extractors.py

"""
Result Extractors

Wrapper functions that integrate with the optimization runner.
These extract optimization metrics from NX Nastran result files.
"""

from pathlib import Path
from typing import Dict, Any
import sys

# Add project root to path
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))

from optimization_engine.result_extractors.op2_extractor_example import (
    extract_max_displacement,
    extract_max_stress,
    extract_mass
)


def mass_extractor(result_path: Path) -> Dict[str, float]:
    """
    Extract mass metrics for optimization.

    Args:
        result_path: Path to .op2 file or directory containing results

    Returns:
        Dict with 'total_mass' and other mass-related metrics
    """
    # If result_path is a directory, find the .op2 file
    if result_path.is_dir():
        op2_files = list(result_path.glob("*.op2"))
        if not op2_files:
            raise FileNotFoundError(f"No .op2 files found in {result_path}")
        op2_path = op2_files[0]  # Use first .op2 file
    else:
        op2_path = result_path

    if not op2_path.exists():
        raise FileNotFoundError(f"Result file not found: {op2_path}")

    result = extract_mass(op2_path)

    # Ensure total_mass key exists
    if 'total_mass' not in result or result['total_mass'] is None:
        raise ValueError(f"Could not extract total_mass from {op2_path}")

    return result


def stress_extractor(result_path: Path) -> Dict[str, float]:
    """
    Extract stress metrics for optimization.

    Args:
        result_path: Path to .op2 file or directory containing results

    Returns:
        Dict with 'max_von_mises' and other stress metrics
    """
    # If result_path is a directory, find the .op2 file
    if result_path.is_dir():
        op2_files = list(result_path.glob("*.op2"))
        if not op2_files:
            raise FileNotFoundError(f"No .op2 files found in {result_path}")
        op2_path = op2_files[0]
    else:
        op2_path = result_path

    if not op2_path.exists():
        raise FileNotFoundError(f"Result file not found: {op2_path}")

    result = extract_max_stress(op2_path, stress_type='von_mises')

    # Ensure max_von_mises key exists
    if 'max_stress' in result:
        result['max_von_mises'] = result['max_stress']

    if 'max_von_mises' not in result or result['max_von_mises'] is None:
        raise ValueError(f"Could not extract max_von_mises from {op2_path}")

    return result


def displacement_extractor(result_path: Path) -> Dict[str, float]:
    """
    Extract displacement metrics for optimization.

    Args:
        result_path: Path to .op2 file or directory containing results

    Returns:
        Dict with 'max_displacement' and other displacement metrics
    """
    # If result_path is a directory, find the .op2 file
    if result_path.is_dir():
        op2_files = list(result_path.glob("*.op2"))
        if not op2_files:
            raise FileNotFoundError(f"No .op2 files found in {result_path}")
        op2_path = op2_files[0]
    else:
        op2_path = result_path

    if not op2_path.exists():
        raise FileNotFoundError(f"Result file not found: {op2_path}")

    result = extract_max_displacement(op2_path)

    # Ensure max_displacement key exists
    if 'max_displacement' not in result or result['max_displacement'] is None:
        raise ValueError(f"Could not extract max_displacement from {op2_path}")

    return result


def volume_extractor(result_path: Path) -> Dict[str, float]:
    """
    Extract volume metrics for optimization.

    Note: Volume is often not directly in OP2 files.
    This is a placeholder that could be extended to:
    - Calculate from mass and density
    - Extract from .f06 file
    - Query from NX model directly

    Args:
        result_path: Path to result files

    Returns:
        Dict with 'total_volume'
    """
    # For now, estimate from mass (would need material density)
    # This is a placeholder implementation
    mass_result = mass_extractor(result_path)

    # Assuming steel density ~7850 kg/m^3 = 7.85e-6 kg/mm^3
    # volume (mm^3) = mass (kg) / density (kg/mm^3)
    assumed_density = 7.85e-6  # kg/mm^3

    if mass_result['total_mass']:
        total_volume = mass_result['total_mass'] / assumed_density
    else:
        total_volume = None

    return {
        'total_volume': total_volume,
        'note': 'Volume estimated from mass using assumed density'
    }


# Registry of all available extractors
EXTRACTOR_REGISTRY = {
    'mass_extractor': mass_extractor,
    'stress_extractor': stress_extractor,
    'displacement_extractor': displacement_extractor,
    'volume_extractor': volume_extractor
}


def get_extractor(extractor_name: str):
    """
    Get an extractor function by name.

    Args:
        extractor_name: Name of the extractor

    Returns:
        Extractor function

    Raises:
        ValueError: If extractor not found
    """
    if extractor_name not in EXTRACTOR_REGISTRY:
        available = ', '.join(EXTRACTOR_REGISTRY.keys())
        raise ValueError(f"Unknown extractor: {extractor_name}. Available: {available}")

    return EXTRACTOR_REGISTRY[extractor_name]


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

    if len(sys.argv) < 2:
        print("Usage: python extractors.py <path_to_op2_file>")
        print("\nAvailable extractors:")
        for name in EXTRACTOR_REGISTRY.keys():
            print(f"  - {name}")
        sys.exit(1)

    result_path = Path(sys.argv[1])

    print("="*60)
    print("TESTING ALL EXTRACTORS")
    print("="*60)

    for extractor_name, extractor_func in EXTRACTOR_REGISTRY.items():
        print(f"\n{extractor_name}:")
        try:
            result = extractor_func(result_path)
            for key, value in result.items():
                print(f"  {key}: {value}")
        except Exception as e:
            print(f"  Error: {e}")