"""
Extract SPC Forces (Reaction Forces) from Structural Analysis
=============================================================

Extracts single-point constraint (SPC) forces from OP2 files.
These are the reaction forces at boundary conditions.

Pattern: spc_forces
Node Type: Constrained grid points
Result Type: reaction force
API: model.spc_forces[subcase]

Phase 2 Task 2.5 - NX Open Automation Roadmap
"""

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


def extract_spc_forces(
    op2_file: Path,
    subcase: int = 1,
    component: Literal['total', 'fx', 'fy', 'fz', 'mx', 'my', 'mz', 'force', 'moment'] = 'total'
) -> Dict[str, Any]:
    """
    Extract SPC (reaction) forces from boundary conditions.

    SPC forces are the reaction forces at constrained nodes. They balance
    the applied loads and indicate load path through the structure.

    Args:
        op2_file: Path to .op2 file
        subcase: Subcase ID (default 1)
        component: Which component(s) to return:
            - 'total': Resultant force magnitude (sqrt(fx^2+fy^2+fz^2))
            - 'fx', 'fy', 'fz': Individual force components
            - 'mx', 'my', 'mz': Individual moment components
            - 'force': Vector sum of forces only
            - 'moment': Vector sum of moments only

    Returns:
        dict: {
            'total_reaction': Total reaction force magnitude,
            'max_reaction': Maximum nodal reaction,
            'max_reaction_node': Node ID with max reaction,
            'sum_fx': Sum of Fx at all nodes,
            'sum_fy': Sum of Fy at all nodes,
            'sum_fz': Sum of Fz at all nodes,
            'sum_mx': Sum of Mx at all nodes,
            'sum_my': Sum of My at all nodes,
            'sum_mz': Sum of Mz at all nodes,
            'node_reactions': Dict of {node_id: [fx,fy,fz,mx,my,mz]},
            'num_constrained_nodes': Number of nodes with SPCs,
            'subcase': Subcase ID,
            'units': 'N, N-mm (model units)'
        }

    Example:
        >>> result = extract_spc_forces('model.op2')
        >>> print(f"Total reaction: {result['total_reaction']:.2f} N")
        >>> print(f"Sum Fz: {result['sum_fz']:.2f} N")
    """
    model = OP2()
    model.read_op2(str(op2_file))

    # Check for SPC forces
    spc_dict = None

    # Try op2_results container first
    if hasattr(model, 'op2_results') and hasattr(model.op2_results, 'force'):
        force_container = model.op2_results.force
        if hasattr(force_container, 'spc_forces'):
            spc_dict = force_container.spc_forces

    # Fallback to direct model attribute
    if spc_dict is None and hasattr(model, 'spc_forces') and model.spc_forces:
        spc_dict = model.spc_forces

    if spc_dict is None or not spc_dict:
        raise ValueError(
            "No SPC forces found in OP2 file. "
            "Ensure SPCFORCE=ALL is in the Nastran case control."
        )

    # Get subcase
    available_subcases = list(spc_dict.keys())
    if subcase in available_subcases:
        actual_subcase = subcase
    else:
        actual_subcase = available_subcases[0]

    spc_result = spc_dict[actual_subcase]

    # Extract data
    # SPC forces: data[itime, inode, 6] where columns are [fx, fy, fz, mx, my, mz]
    itime = 0

    if hasattr(spc_result, 'data'):
        data = spc_result.data[itime]  # Shape: (nnodes, 6)
        node_ids = spc_result.node_gridtype[:, 0] if hasattr(spc_result, 'node_gridtype') else np.arange(len(data))
    else:
        raise ValueError("Unexpected SPC forces data format")

    # Force components (columns 0-2)
    fx = data[:, 0]
    fy = data[:, 1]
    fz = data[:, 2]

    # Moment components (columns 3-5)
    mx = data[:, 3]
    my = data[:, 4]
    mz = data[:, 5]

    # Resultant force at each node
    force_mag = np.sqrt(fx**2 + fy**2 + fz**2)
    moment_mag = np.sqrt(mx**2 + my**2 + mz**2)

    # Total resultant (force + moment contribution)
    # Note: Forces and moments have different units, so total is just force magnitude
    total_mag = force_mag

    # Statistics
    max_idx = np.argmax(total_mag)
    max_reaction = float(total_mag[max_idx])
    max_node = int(node_ids[max_idx])

    # Sum of components (should balance applied loads in static analysis)
    sum_fx = float(np.sum(fx))
    sum_fy = float(np.sum(fy))
    sum_fz = float(np.sum(fz))
    sum_mx = float(np.sum(mx))
    sum_my = float(np.sum(my))
    sum_mz = float(np.sum(mz))

    # Total reaction force magnitude
    total_reaction = float(np.sqrt(sum_fx**2 + sum_fy**2 + sum_fz**2))

    # Per-node reactions dictionary
    node_reactions = {}
    for i, nid in enumerate(node_ids):
        node_reactions[int(nid)] = [
            float(fx[i]), float(fy[i]), float(fz[i]),
            float(mx[i]), float(my[i]), float(mz[i])
        ]

    # Component-specific return values
    component_result = None
    if component == 'fx':
        component_result = float(np.max(np.abs(fx)))
    elif component == 'fy':
        component_result = float(np.max(np.abs(fy)))
    elif component == 'fz':
        component_result = float(np.max(np.abs(fz)))
    elif component == 'mx':
        component_result = float(np.max(np.abs(mx)))
    elif component == 'my':
        component_result = float(np.max(np.abs(my)))
    elif component == 'mz':
        component_result = float(np.max(np.abs(mz)))
    elif component == 'force':
        component_result = float(np.max(force_mag))
    elif component == 'moment':
        component_result = float(np.max(moment_mag))
    elif component == 'total':
        component_result = total_reaction

    return {
        'total_reaction': total_reaction,
        'max_reaction': max_reaction,
        'max_reaction_node': max_node,
        'component_max': component_result,
        'component': component,
        'sum_fx': sum_fx,
        'sum_fy': sum_fy,
        'sum_fz': sum_fz,
        'sum_mx': sum_mx,
        'sum_my': sum_my,
        'sum_mz': sum_mz,
        'max_fx': float(np.max(np.abs(fx))),
        'max_fy': float(np.max(np.abs(fy))),
        'max_fz': float(np.max(np.abs(fz))),
        'max_mx': float(np.max(np.abs(mx))),
        'max_my': float(np.max(np.abs(my))),
        'max_mz': float(np.max(np.abs(mz))),
        'node_reactions': node_reactions,
        'num_constrained_nodes': len(node_ids),
        'subcase': actual_subcase,
        'units': 'N, N-mm (model units)',
    }


def extract_total_reaction_force(
    op2_file: Path,
    subcase: int = 1
) -> float:
    """
    Convenience function to extract total reaction force magnitude.

    Args:
        op2_file: Path to .op2 file
        subcase: Subcase ID

    Returns:
        Total reaction force magnitude (N)
    """
    result = extract_spc_forces(op2_file, subcase)
    return result['total_reaction']


def extract_reaction_component(
    op2_file: Path,
    component: str = 'fz',
    subcase: int = 1
) -> float:
    """
    Extract maximum absolute value of a specific reaction component.

    Args:
        op2_file: Path to .op2 file
        component: 'fx', 'fy', 'fz', 'mx', 'my', 'mz'
        subcase: Subcase ID

    Returns:
        Maximum absolute value of the specified component
    """
    result = extract_spc_forces(op2_file, subcase, component)
    return result['component_max']


def check_force_equilibrium(
    op2_file: Path,
    applied_load: Optional[Dict[str, float]] = None,
    tolerance: float = 1.0
) -> Dict[str, Any]:
    """
    Check if reaction forces balance applied loads (equilibrium check).

    In a valid static analysis, sum of reactions should equal applied loads.

    Args:
        op2_file: Path to .op2 file
        applied_load: Optional dict of applied loads {'fx': N, 'fy': N, 'fz': N}
        tolerance: Tolerance for equilibrium check (default 1.0 N)

    Returns:
        dict: {
            'in_equilibrium': Boolean,
            'reaction_sum': [fx, fy, fz],
            'imbalance': [dx, dy, dz] (if applied_load provided),
            'max_imbalance': Maximum component imbalance
        }
    """
    result = extract_spc_forces(op2_file)

    reaction_sum = [result['sum_fx'], result['sum_fy'], result['sum_fz']]

    equilibrium_result = {
        'reaction_sum': reaction_sum,
        'moment_sum': [result['sum_mx'], result['sum_my'], result['sum_mz']],
    }

    if applied_load:
        applied = [
            applied_load.get('fx', 0.0),
            applied_load.get('fy', 0.0),
            applied_load.get('fz', 0.0)
        ]
        # Reactions should be opposite to applied loads
        imbalance = [
            reaction_sum[0] + applied[0],
            reaction_sum[1] + applied[1],
            reaction_sum[2] + applied[2]
        ]
        max_imbalance = max(abs(i) for i in imbalance)
        equilibrium_result['applied_load'] = applied
        equilibrium_result['imbalance'] = imbalance
        equilibrium_result['max_imbalance'] = max_imbalance
        equilibrium_result['in_equilibrium'] = max_imbalance <= tolerance
    else:
        # Without applied load, just check if reactions are non-zero
        equilibrium_result['total_reaction'] = result['total_reaction']
        equilibrium_result['in_equilibrium'] = True  # Can't check without applied load

    return equilibrium_result


if __name__ == '__main__':
    import sys
    if len(sys.argv) > 1:
        op2_file = Path(sys.argv[1])

        try:
            result = extract_spc_forces(op2_file)

            print(f"\nSPC Forces (Reaction Forces):")
            print(f"  Total reaction: {result['total_reaction']:.2f} N")
            print(f"  Max nodal reaction: {result['max_reaction']:.2f} N")
            print(f"    Node: {result['max_reaction_node']}")
            print(f"\n  Force sums (should balance applied loads):")
            print(f"    ΣFx: {result['sum_fx']:.2f} N")
            print(f"    ΣFy: {result['sum_fy']:.2f} N")
            print(f"    ΣFz: {result['sum_fz']:.2f} N")
            print(f"\n  Moment sums:")
            print(f"    ΣMx: {result['sum_mx']:.2f} N-mm")
            print(f"    ΣMy: {result['sum_my']:.2f} N-mm")
            print(f"    ΣMz: {result['sum_mz']:.2f} N-mm")
            print(f"\n  Constrained nodes: {result['num_constrained_nodes']}")

            # Show a few node reactions
            print(f"\n  Sample node reactions:")
            for i, (nid, forces) in enumerate(result['node_reactions'].items()):
                if i >= 3:
                    print(f"    ... ({result['num_constrained_nodes'] - 3} more)")
                    break
                print(f"    Node {nid}: Fx={forces[0]:.2f}, Fy={forces[1]:.2f}, Fz={forces[2]:.2f}")

        except ValueError as e:
            print(f"Error: {e}")
    else:
        print(f"Usage: python {sys.argv[0]} <op2_file>")