Atomizer/optimization_engine/extractors/op2_extractor.py

"""
Generic OP2 Extractor
====================

Reusable extractor for NX Nastran OP2 files using pyNastran.
Extracts mass properties, forces, displacements, stresses, etc.

Usage:
    extractor = OP2Extractor(op2_file="model.op2")
    mass = extractor.extract_mass()
    forces = extractor.extract_grid_point_forces()
"""

from pathlib import Path
from typing import Dict, Any, Optional, List
import numpy as np

try:
    from pyNastran.op2.op2 import read_op2
except ImportError:
    raise ImportError("pyNastran is required. Install with: pip install pyNastran")


class OP2Extractor:
    """
    Generic extractor for Nastran OP2 files.

    Supports:
    - Mass properties
    - Grid point forces
    - Displacements
    - Stresses
    - Strains
    - Element forces
    """

    def __init__(self, op2_file: str):
        """
        Args:
            op2_file: Path to .op2 file
        """
        self.op2_file = Path(op2_file)
        self._op2_model = None

    def _load_op2(self):
        """Lazy load OP2 file"""
        if self._op2_model is None:
            if not self.op2_file.exists():
                raise FileNotFoundError(f"OP2 file not found: {self.op2_file}")
            self._op2_model = read_op2(str(self.op2_file), debug=False)
        return self._op2_model

    def extract_mass(self, subcase_id: Optional[int] = None) -> Dict[str, Any]:
        """
        Extract mass properties from OP2.

        Returns:
            dict: {
                'mass_kg': total mass in kg,
                'mass_g': total mass in grams,
                'cg': [x, y, z] center of gravity,
                'inertia': 3x3 inertia matrix
            }
        """
        op2 = self._load_op2()

        # Get grid point weight (mass properties)
        if not hasattr(op2, 'grid_point_weight') or not op2.grid_point_weight:
            raise ValueError("No mass properties found in OP2 file")

        gpw = op2.grid_point_weight

        # Mass is typically in the first element of MO matrix (reference point mass)
        # OP2 stores mass in ton, mm, sec units typically
        mass_matrix = gpw.MO[0, 0] if hasattr(gpw, 'MO') else None

        # Get reference point
        if hasattr(gpw, 'reference_point') and gpw.reference_point:
            ref_point = gpw.reference_point
        else:
            ref_point = 0

        # Extract mass (convert based on units)
        # Nastran default: ton-mm-sec → need to convert to kg
        if mass_matrix is not None:
            mass_ton = mass_matrix
            mass_kg = mass_ton * 1000.0  # 1 ton = 1000 kg
        else:
            raise ValueError("Could not extract mass from OP2")

        # Extract CG if available
        cg = [0.0, 0.0, 0.0]
        if hasattr(gpw, 'cg'):
            cg = gpw.cg.tolist() if hasattr(gpw.cg, 'tolist') else list(gpw.cg)

        return {
            'mass_kg': mass_kg,
            'mass_g': mass_kg * 1000.0,
            'mass_ton': mass_ton,
            'cg': cg,
            'reference_point': ref_point,
            'units': 'ton-mm-sec (converted to kg)',
        }

    def extract_grid_point_forces(
        self,
        subcase_id: Optional[int] = None,
        component: str = "total"  # total, fx, fy, fz, mx, my, mz
    ) -> Dict[str, Any]:
        """
        Extract grid point forces from OP2.

        Args:
            subcase_id: Subcase ID (if None, uses first available)
            component: Force component to extract

        Returns:
            dict: {
                'force': resultant force value,
                'all_forces': list of forces at each grid point,
                'max_force': maximum force,
                'total_force': sum of all forces
            }
        """
        op2 = self._load_op2()

        if not hasattr(op2, 'grid_point_forces') or not op2.grid_point_forces:
            raise ValueError("No grid point forces found in OP2 file")

        # Get first subcase if not specified
        if subcase_id is None:
            subcase_id = list(op2.grid_point_forces.keys())[0]

        gpf = op2.grid_point_forces[subcase_id]

        # Extract forces based on component
        # Grid point forces table typically has columns: fx, fy, fz, mx, my, mz
        if component == "total":
            # Calculate resultant force: sqrt(fx^2 + fy^2 + fz^2)
            forces = np.sqrt(gpf.data[:, 0]**2 + gpf.data[:, 1]**2 + gpf.data[:, 2]**2)
        elif component == "fx":
            forces = gpf.data[:, 0]
        elif component == "fy":
            forces = gpf.data[:, 1]
        elif component == "fz":
            forces = gpf.data[:, 2]
        else:
            raise ValueError(f"Unknown component: {component}")

        return {
            'force': float(np.max(np.abs(forces))),
            'all_forces': forces.tolist(),
            'max_force': float(np.max(forces)),
            'min_force': float(np.min(forces)),
            'total_force': float(np.sum(forces)),
            'component': component,
            'subcase_id': subcase_id,
        }

    def extract_applied_loads(self, subcase_id: Optional[int] = None) -> Dict[str, Any]:
        """
        Extract applied loads from OP2 file.

        This attempts to get load vector information if available.
        Note: Not all OP2 files contain this data.

        Returns:
            dict: Load information
        """
        op2 = self._load_op2()

        # Try to get load vectors
        if hasattr(op2, 'load_vectors') and op2.load_vectors:
            if subcase_id is None:
                subcase_id = list(op2.load_vectors.keys())[0]

            lv = op2.load_vectors[subcase_id]
            loads = lv.data

            return {
                'total_load': float(np.sum(np.abs(loads))),
                'max_load': float(np.max(np.abs(loads))),
                'load_resultant': float(np.linalg.norm(loads)),
                'subcase_id': subcase_id,
            }
        else:
            # Fallback: use grid point forces as approximation
            return self.extract_grid_point_forces(subcase_id)


def extract_mass_from_op2(op2_file: str) -> float:
    """
    Convenience function to extract mass in kg.

    Args:
        op2_file: Path to .op2 file

    Returns:
        Mass in kilograms
    """
    extractor = OP2Extractor(op2_file)
    result = extractor.extract_mass()
    return result['mass_kg']


def extract_force_from_op2(
    op2_file: str,
    component: str = "fz"
) -> float:
    """
    Convenience function to extract force component.

    Args:
        op2_file: Path to .op2 file
        component: Force component (fx, fy, fz, or total)

    Returns:
        Force value
    """
    extractor = OP2Extractor(op2_file)
    result = extractor.extract_grid_point_forces(component=component)
    return result['force']


if __name__ == "__main__":
    # Example usage
    import sys
    if len(sys.argv) > 1:
        op2_file = sys.argv[1]
        extractor = OP2Extractor(op2_file)

        # Extract mass
        mass_result = extractor.extract_mass()
        print(f"Mass: {mass_result['mass_kg']:.6f} kg")
        print(f"CG: {mass_result['cg']}")

        # Extract forces
        try:
            force_result = extractor.extract_grid_point_forces(component="fz")
            print(f"Max Fz: {force_result['force']:.2f} N")
        except ValueError as e:
            print(f"Forces not available: {e}")