"""
Example: Result Extraction from OP2 files using pyNastran

This shows how to extract optimization metrics from Nastran OP2 files.
Common metrics:
- Max displacement (for stiffness constraints)
- Max von Mises stress (for strength constraints)
- Mass (for minimization objectives)
"""

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


def extract_max_displacement(op2_path: Path) -> Dict[str, Any]:
    """
    Extract maximum displacement magnitude from OP2 file.

    Args:
        op2_path: Path to .op2 file

    Returns:
        Dictionary with max displacement, node ID, and components
    """
    from pyNastran.op2.op2 import OP2

    op2 = OP2()
    op2.read_op2(str(op2_path))

    # Get first subcase (usually the only one in static analysis)
    subcase_id = list(op2.displacements.keys())[0]
    displacements = op2.displacements[subcase_id]

    # Extract node IDs and displacement data
    node_ids = displacements.node_gridtype[:, 0].astype(int)
    disp_data = displacements.data[0]  # First (and usually only) timestep

    # Calculate magnitude: sqrt(dx^2 + dy^2 + dz^2)
    dx = disp_data[:, 0]
    dy = disp_data[:, 1]
    dz = disp_data[:, 2]
    magnitudes = np.sqrt(dx**2 + dy**2 + dz**2)

    # Find max
    max_idx = np.argmax(magnitudes)
    max_displacement = magnitudes[max_idx]
    max_node_id = node_ids[max_idx]

    return {
        'max_displacement': float(max_displacement),
        'max_node_id': int(max_node_id),
        'dx': float(dx[max_idx]),
        'dy': float(dy[max_idx]),
        'dz': float(dz[max_idx]),
        'units': 'mm',  # NX typically uses mm
        'subcase': subcase_id
    }


def extract_max_stress(op2_path: Path, stress_type: str = 'von_mises') -> Dict[str, Any]:
    """
    Extract maximum stress from OP2 file.

    Args:
        op2_path: Path to .op2 file
        stress_type: 'von_mises' or 'max_principal'

    Returns:
        Dictionary with max stress, element ID, and location
    """
    from pyNastran.op2.op2 import OP2

    op2 = OP2(debug=False)
    op2.read_op2(str(op2_path))

    # Stress can be in different tables depending on element type
    # Common: cquad4_stress, ctria3_stress, ctetra_stress, etc.
    stress_tables = [
        'cquad4_stress',
        'ctria3_stress',
        'ctetra_stress',
        'chexa_stress',
        'cpenta_stress',
        'cbar_stress',
        'cbeam_stress'
    ]

    max_stress_overall = 0.0
    max_element_id = None
    max_element_type = None

    # Try to get stress from different pyNastran API formats
    for table_name in stress_tables:
        stress_table = None

        # Try format 1: Attribute name with dot (e.g., 'stress.chexa_stress')
        # This is used in newer pyNastran versions
        dotted_name = f'stress.{table_name}'
        if hasattr(op2, dotted_name):
            stress_table = getattr(op2, dotted_name)
        # Try format 2: Nested attribute op2.stress.chexa_stress
        elif hasattr(op2, 'stress') and hasattr(op2.stress, table_name):
            stress_table = getattr(op2.stress, table_name)
        # Try format 3: Direct attribute op2.chexa_stress (older pyNastran)
        elif hasattr(op2, table_name):
            stress_table = getattr(op2, table_name)

        if stress_table:
            subcase_id = list(stress_table.keys())[0]
            stress_data = stress_table[subcase_id]

            # Extract von Mises stress
            # Note: Structure varies by element type
            element_ids = stress_data.element_node[:, 0].astype(int)

            if stress_type == 'von_mises':
                # For solid elements (CHEXA, CTETRA, CPENTA): von Mises is at index 9
                # For shell elements (CQUAD4, CTRIA3): von Mises is last column (-1)
                if table_name in ['chexa_stress', 'ctetra_stress', 'cpenta_stress']:
                    # Solid elements: data shape is [itime, nnodes, 10]
                    # Index 9 is von_mises [oxx, oyy, ozz, txy, tyz, txz, o1, o2, o3, von_mises]
                    stresses = stress_data.data[0, :, 9]
                else:
                    # Shell elements: von Mises is last column
                    stresses = stress_data.data[0, :, -1]
            else:
                # Max principal stress
                if table_name in ['chexa_stress', 'ctetra_stress', 'cpenta_stress']:
                    stresses = stress_data.data[0, :, 6]  # o1 (max principal)
                else:
                    stresses = stress_data.data[0, :, -2]

            max_stress_in_table = np.max(stresses)
            if max_stress_in_table > max_stress_overall:
                max_stress_overall = max_stress_in_table
                max_idx = np.argmax(stresses)
                max_element_id = element_ids[max_idx]
                max_element_type = table_name.replace('_stress', '')

    # CRITICAL: NX Nastran outputs stress in kPa (mN/mm²), convert to MPa
    # 1 kPa = 0.001 MPa
    max_stress_overall_mpa = max_stress_overall / 1000.0

    return {
        'max_stress': float(max_stress_overall_mpa),
        'stress_type': stress_type,
        'element_id': int(max_element_id) if max_element_id else None,
        'element_type': max_element_type,
        'units': 'MPa',
    }


def extract_mass(op2_path: Path) -> Dict[str, Any]:
    """
    Extract total mass from OP2 file.

    Args:
        op2_path: Path to .op2 file

    Returns:
        Dictionary with mass and center of gravity
    """
    from pyNastran.op2.op2 import OP2

    op2 = OP2()
    op2.read_op2(str(op2_path))

    # Mass is in grid_point_weight table
    if hasattr(op2, 'grid_point_weight') and op2.grid_point_weight:
        mass_data = op2.grid_point_weight

        # Total mass
        total_mass = mass_data.mass.sum()

        # Center of gravity
        cg = mass_data.cg

        return {
            'total_mass': float(total_mass),
            'cg_x': float(cg[0]),
            'cg_y': float(cg[1]),
            'cg_z': float(cg[2]),
            'units': 'kg'
        }
    else:
        # Fallback: Mass not directly available
        return {
            'total_mass': None,
            'note': 'Mass data not found in OP2 file. Ensure PARAM,GRDPNT,0 is in Nastran deck'
        }


# Combined extraction function for optimization
def extract_all_results(op2_path: Path) -> Dict[str, Any]:
    """
    Extract all common optimization metrics from OP2 file.

    Args:
        op2_path: Path to .op2 file

    Returns:
        Dictionary with all results
    """
    results = {
        'op2_file': str(op2_path),
        'status': 'success'
    }

    try:
        results['displacement'] = extract_max_displacement(op2_path)
    except Exception as e:
        results['displacement'] = {'error': str(e)}

    try:
        results['stress'] = extract_max_stress(op2_path)
    except Exception as e:
        results['stress'] = {'error': str(e)}

    try:
        results['mass'] = extract_mass(op2_path)
    except Exception as e:
        results['mass'] = {'error': str(e)}

    return results


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

    if len(sys.argv) < 2:
        print("Usage: python op2_extractor_example.py <path_to_op2_file>")
        sys.exit(1)

    op2_path = Path(sys.argv[1])

    if not op2_path.exists():
        print(f"Error: File not found: {op2_path}")
        sys.exit(1)

    print(f"Extracting results from: {op2_path}")
    print("=" * 60)

    results = extract_all_results(op2_path)

    print("\nResults:")
    print(json.dumps(results, indent=2))

    # Summary
    print("\n" + "=" * 60)
    print("SUMMARY:")
    if 'displacement' in results and 'max_displacement' in results['displacement']:
        disp = results['displacement']
        print(f"  Max Displacement: {disp['max_displacement']:.6f} {disp['units']} at node {disp['max_node_id']}")

    if 'stress' in results and 'max_stress' in results['stress']:
        stress = results['stress']
        print(f"  Max {stress['stress_type']}: {stress['max_stress']:.2f} {stress['units']} in element {stress['element_id']}")

    if 'mass' in results and 'total_mass' in results['mass'] and results['mass']['total_mass']:
        mass = results['mass']
        print(f"  Total Mass: {mass['total_mass']:.6f} {mass['units']}")