"""
Extract maximum von Mises stress from structural analysis.

Supports all solid element types (CTETRA, CHEXA, CPENTA, CPYRAM) and
shell elements (CQUAD4, CTRIA3).

Unit Note: NX Nastran in kg-mm-s outputs stress in kPa. This extractor
converts to MPa (divide by 1000) for engineering use.
"""

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


def extract_solid_stress(
    op2_file: Path,
    subcase: int = 1,
    element_type: Optional[str] = None,
    convert_to_mpa: bool = True,
) -> Dict[str, Any]:
    """
    Extract maximum von Mises stress from solid elements.

    Args:
        op2_file: Path to OP2 results file
        subcase: Subcase ID (default 1)
        element_type: Specific element type to check ('ctetra', 'chexa', etc.)
                     If None, checks ALL solid element types and returns max.
        convert_to_mpa: If True, divide by 1000 to convert kPa to MPa (default True)

    Returns:
        dict with 'max_von_mises' (in MPa if convert_to_mpa=True),
        'max_stress_element', and 'element_type'
    """
    model = OP2(debug=False, log=None)
    model.read_op2(str(op2_file))

    # All solid element types to check
    solid_element_types = ["ctetra", "chexa", "cpenta", "cpyram"]
    shell_element_types = ["cquad4", "ctria3"]

    # If specific element type requested, only check that one
    if element_type:
        element_types_to_check = [element_type.lower()]
    else:
        # Check all solid types by default
        element_types_to_check = solid_element_types

    if not hasattr(model, "op2_results") or not hasattr(model.op2_results, "stress"):
        raise ValueError("No stress results in OP2 file")

    stress_container = model.op2_results.stress

    # Find max stress across all requested element types
    max_stress = 0.0
    max_stress_elem = 0
    max_stress_type = None

    for elem_type in element_types_to_check:
        stress_attr = f"{elem_type}_stress"

        if not hasattr(stress_container, stress_attr):
            continue

        stress_dict = getattr(stress_container, stress_attr)
        if not stress_dict:
            continue

        # Get subcase
        available_subcases = list(stress_dict.keys())
        if not available_subcases:
            continue

        actual_subcase = subcase if subcase in available_subcases else available_subcases[0]
        stress = stress_dict[actual_subcase]

        if not stress.is_von_mises:
            continue

        # Determine von Mises column
        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:
            von_mises_col = ncols - 1

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

        if elem_max > max_stress:
            max_stress = elem_max
            element_ids = [eid for (eid, node) in stress.element_node]
            max_stress_elem = int(element_ids[np.argmax(von_mises)])
            max_stress_type = elem_type.upper()

    if max_stress_type is None:
        raise ValueError(f"No stress results found for element types: {element_types_to_check}")

    # Convert from kPa to MPa (NX kg-mm-s unit system outputs kPa)
    if convert_to_mpa:
        max_stress = max_stress / 1000.0

    return {
        "max_von_mises": max_stress,
        "max_stress_element": max_stress_elem,
        "element_type": max_stress_type,
        "units": "MPa" if convert_to_mpa else "kPa",
    }


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