"""
Extract von Mises stress from solid and shell elements
Auto-generated by Atomizer Phase 3 - pyNastran Research Agent

Pattern: element_stress
Element Types: CTETRA, CHEXA, CPENTA (solids), CQUAD4, CTRIA3 (shells)
Result Type: stress
API: model.ctetra_stress[subcase], model.cquad4_stress[subcase], etc.
"""

from pathlib import Path
from typing import Dict, Any
import numpy as np
from pyNastran.op2.op2 import OP2


def extract_solid_stress(op2_file: Path, subcase: int = 1, element_type: str = 'auto'):
    """
    Extract stress from solid or shell elements.

    Args:
        op2_file: Path to OP2 file
        subcase: Subcase number (default 1)
        element_type: Element type ('ctetra', 'chexa', 'cquad4', 'ctria3', or 'auto')
                     'auto' will detect available element types

    Returns:
        Dict with max von Mises stress and element info
    """
    from pyNastran.op2.op2 import OP2
    import numpy as np

    model = OP2()
    model.read_op2(str(op2_file))

    # Auto-detect element type if requested
    if element_type == 'auto':
        # Try common element types in order
        # pyNastran uses "stress.{element}_stress" as attribute names (with dot in name!)
        possible_types = ['cquad4', 'ctria3', 'ctetra', 'chexa', 'cpenta']
        element_type = None

        for elem_type in possible_types:
            stress_attr = f"stress.{elem_type}_stress"
            try:
                stress_dict = getattr(model, stress_attr, None)
                if isinstance(stress_dict, dict) and subcase in stress_dict:
                    element_type = elem_type
                    break
            except AttributeError:
                continue

        if element_type is None:
            raise ValueError(f"No stress results found in OP2 for subcase {subcase}")

    # Get stress object for element type
    # pyNastran stores stress as "stress.{element}_stress" (e.g., "stress.cquad4_stress")
    stress_attr = f"stress.{element_type}_stress"

    try:
        stress_dict = getattr(model, stress_attr)
    except AttributeError:
        raise ValueError(f"No {element_type} stress results in OP2")

    if not isinstance(stress_dict, dict) or subcase not in stress_dict:
        raise ValueError(f"Subcase {subcase} not found in {element_type} stress results")

    stress = stress_dict[subcase]
    itime = 0

    # Extract von Mises
    # For CQUAD4/CTRIA3 (shells): data shape is [ntimes, nelements, 8]
    # Columns: [fiber_distance, oxx, oyy, txy, angle, omax, omin, von_mises]
    # Column 7 (0-indexed) is von Mises
    #
    # For CTETRA/CHEXA (solids): Column 9 is von Mises

    if element_type in ['cquad4', 'ctria3']:
        # Shell elements - von Mises is at column 7
        von_mises = stress.data[itime, :, 7]
    else:
        # Solid elements - von Mises is at column 9
        von_mises = stress.data[itime, :, 9]

    # Raw stress values from OP2 (in internal Nastran units)
    max_stress_raw = float(np.max(von_mises))
    min_stress_raw = float(np.min(von_mises))
    avg_stress_raw = float(np.mean(von_mises))

    # Convert to MPa
    # For MN-MM unit system (UNITSYS=MN-MM), Nastran outputs stress with implied decimal
    # The raw value needs to be divided by 1000 to get MPa
    # Example: 131507.1875 (raw) = 131.507 MPa
    max_stress_mpa = max_stress_raw / 1000.0
    min_stress_mpa = min_stress_raw / 1000.0
    avg_stress_mpa = avg_stress_raw / 1000.0

    # Get element info
    if hasattr(stress, 'element_node'):
        element_ids = stress.element_node[:, 0]  # First column is element ID
    elif hasattr(stress, 'element'):
        element_ids = stress.element
    else:
        element_ids = np.arange(len(von_mises))

    max_stress_elem = int(element_ids[np.argmax(von_mises)])

    return {
        'max_von_mises': max_stress_mpa,
        'min_von_mises': min_stress_mpa,
        'avg_von_mises': avg_stress_mpa,
        'max_stress_element': max_stress_elem,
        'element_type': element_type,
        'num_elements': len(von_mises),
        'units': 'MPa'
    }


if __name__ == '__main__':
    # Example usage
    import sys
    if len(sys.argv) > 1:
        op2_file = Path(sys.argv[1])
        result = extract_solid_stress(op2_file)
        print(f"Extraction result: {result}")
    else:
        print("Usage: python {sys.argv[0]} <op2_file>")