"""
Extract per-element von Mises stress within each sandbox region.

Reads from the Nastran OP2 (stress per element) + NX FEM file (node coordinates),
then filters to elements whose centroids fall inside the sandbox boundary polygon.

This gives the spatial stress field in sandbox 2D coordinates — the same schema
that solve_and_extract.py would produce via NXOpen once wired:

    {
        "nodes_xy":      [[x, y], ...],   # element centroids (mm)
        "stress_values": [...],            # von Mises per element (MPa)
        "n_elements":    int,
    }

The sandbox polygon filter (Shapely point-in-polygon) is what maps "whole plate"
OP2 stress back to each individual sandbox region.
"""

from __future__ import annotations

from pathlib import Path
from typing import Dict, Any, List

import numpy as np
from shapely.geometry import Polygon, Point


def extract_sandbox_stress_field(
    op2_file: Path,
    fem_file: Path,
    sandbox_geometry: dict,
    subcase: int = 1,
    convert_to_mpa: bool = True,
) -> Dict[str, Any]:
    """
    Extract Von Mises stress field for one sandbox region.

    Args:
        op2_file:         Nastran OP2 results file
        fem_file:         NX FEM file (BDF format) — for node coordinates
        sandbox_geometry: Geometry dict with 'outer_boundary' polygon
        subcase:          Nastran subcase ID (default 1)
        convert_to_mpa:   Divide by 1000 (NX kg-mm-s outputs kPa → MPa)

    Returns:
        Dict with 'nodes_xy', 'stress_values', 'n_elements'
        Returns empty result (n_elements=0) on any failure.
    """
    _empty = {"nodes_xy": [], "stress_values": [], "n_elements": 0}

    try:
        from pyNastran.op2.op2 import OP2
        from pyNastran.bdf.bdf import BDF
    except ImportError:
        print("  [StressField] pyNastran not available — skipping")
        return _empty

    # ── 1. Read FEM for node positions ───────────────────────────────────────
    try:
        bdf = BDF(debug=False)
        bdf.read_bdf(str(fem_file), xref=True)
    except Exception as e:
        print(f"  [StressField] FEM read failed: {e}")
        return _empty

    node_pos: Dict[int, np.ndarray] = {}
    for nid, node in bdf.nodes.items():
        try:
            node_pos[nid] = node.get_position()  # [x, y, z] in mm
        except Exception:
            pass

    # ── 2. Read OP2 for per-element stress ────────────────────────────────────
    try:
        model = OP2(debug=False, log=None)
        model.read_op2(str(op2_file))
    except Exception as e:
        print(f"  [StressField] OP2 read failed: {e}")
        return _empty

    if not hasattr(model, "op2_results") or not hasattr(model.op2_results, "stress"):
        print("  [StressField] No stress results in OP2")
        return _empty

    stress_container = model.op2_results.stress

    # Accumulate per-element von Mises (average over integration points)
    eid_vm_lists: Dict[int, List[float]] = {}

    for etype in ("ctetra", "chexa", "cpenta", "cpyram"):
        attr = f"{etype}_stress"
        if not hasattr(stress_container, attr):
            continue
        stress_dict = getattr(stress_container, attr)
        if not stress_dict:
            continue

        available = list(stress_dict.keys())
        actual_sub = subcase if subcase in available else available[0]
        stress = stress_dict[actual_sub]

        if not stress.is_von_mises:
            continue

        ncols = stress.data.shape[2]
        vm_col = 9 if ncols >= 10 else (7 if ncols == 8 else ncols - 1)

        itime = 0
        for row_idx, (eid, _nid) in enumerate(stress.element_node):
            vm = float(stress.data[itime, row_idx, vm_col])
            eid_vm_lists.setdefault(eid, []).append(vm)

    if not eid_vm_lists:
        print("  [StressField] No solid element stress found in OP2")
        return _empty

    # Average integration points → one value per element
    eid_vm = {eid: float(np.mean(vals)) for eid, vals in eid_vm_lists.items()}

    # ── 3. Compute element centroids + filter to sandbox ─────────────────────
    sandbox_polygon = Polygon(sandbox_geometry["outer_boundary"])

    nodes_xy: List[List[float]] = []
    stress_vals: List[float] = []

    for eid, vm in eid_vm.items():
        if eid not in bdf.elements:
            continue
        try:
            nids = bdf.elements[eid].node_ids
        except Exception:
            continue

        pts = [node_pos[nid] for nid in nids if nid in node_pos]
        if not pts:
            continue

        centroid = np.mean(pts, axis=0)       # [x, y, z] 3D
        x, y = float(centroid[0]), float(centroid[1])  # flat plate → Z constant

        if not sandbox_polygon.contains(Point(x, y)):
            continue

        # Convert kPa → MPa (NX kg-mm-s unit system)
        val = vm / 1000.0 if convert_to_mpa else vm

        nodes_xy.append([x, y])
        stress_vals.append(val)

    n = len(stress_vals)
    if n > 0:
        print(f"  [StressField] sandbox '{sandbox_geometry.get('sandbox_id', '?')}': "
              f"{n} elements  max={max(stress_vals):.1f} MPa")
    else:
        print(f"  [StressField] sandbox '{sandbox_geometry.get('sandbox_id', '?')}': "
              f"0 elements in polygon (check coordinate frame)")

    return {
        "nodes_xy":      nodes_xy,
        "stress_values": stress_vals,
        "n_elements":    n,
    }