"""
Extract maximum von Mises stress from structural analysis
Auto-generated by Atomizer Phase 3 - pyNastran Research Agent

Pattern: solid_stress
Element Type: CTETRA
Result Type: stress
API: model.ctetra_stress[subcase] or model.chexa_stress[subcase]
"""

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 = 'ctetra'):
    """Extract stress from solid elements."""
    from pyNastran.op2.op2 import OP2
    import numpy as np

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

    # Get stress object for element type
    # Different element types have different stress attributes
    stress_attr_map = {
        'ctetra': 'ctetra_stress',
        'chexa': 'chexa_stress',
        'cquad4': 'cquad4_stress',
        'ctria3': 'ctria3_stress'
    }

    stress_attr = stress_attr_map.get(element_type.lower())
    if not stress_attr:
        raise ValueError(f"Unknown element type: {element_type}")

    # Access stress through op2_results container
    # pyNastran structure: model.op2_results.stress.cquad4_stress[subcase]
    stress_dict = None

    if hasattr(model, 'op2_results') and hasattr(model.op2_results, 'stress'):
        stress_container = model.op2_results.stress
        if hasattr(stress_container, stress_attr):
            stress_dict = getattr(stress_container, stress_attr)

    if stress_dict is None:
        raise ValueError(f"No {element_type} stress results in OP2. Available attributes: {[a for a in dir(model) if 'stress' in a.lower()]}")

    # stress_dict is a dictionary with subcase IDs as keys
    available_subcases = list(stress_dict.keys())
    if not available_subcases:
        raise ValueError(f"No stress data found in OP2 file")

    # Use the specified subcase or first available
    if subcase in available_subcases:
        actual_subcase = subcase
    else:
        actual_subcase = available_subcases[0]

    stress = stress_dict[actual_subcase]

    itime = 0

    # Extract von Mises if available
    if stress.is_von_mises:  # Property, not method
        # Different element types have von Mises at different column indices
        # Shell elements (CQUAD4, CTRIA3): 8 columns, von Mises at column 7
        # Solid elements (CTETRA, CHEXA): 10 columns, von Mises at column 9
        ncols = stress.data.shape[2]

        if ncols == 8:
            # Shell elements - von Mises is last column
            von_mises_col = 7
        elif ncols >= 10:
            # Solid elements - von Mises is column 9
            von_mises_col = 9
        else:
            # Unknown format, try last column
            von_mises_col = ncols - 1

        von_mises = stress.data[itime, :, von_mises_col]
        max_stress = float(np.max(von_mises))

        # Get element info
        element_ids = [eid for (eid, node) in stress.element_node]
        max_stress_elem = element_ids[np.argmax(von_mises)]

        return {
            'max_von_mises': max_stress,
            'max_stress_element': int(max_stress_elem)
        }
    else:
        raise ValueError("von Mises stress not available")


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>")