Atomizer/optimization_engine/extractors/extract_stress_field_2d.py

"""
Extract a 2D von Mises stress field from OP2 results, projected onto the sandbox plane.

Works for both:
  - 3D solid meshes (CHEXA, CTETRA, CPENTA): averages stress through thickness
  - 2D shell meshes (CQUAD4, CTRIA3): directly maps to plane

The returned field is in the sandbox 2D coordinate system (u, v) matching the
geometry_sandbox_N.json coordinate space — ready to feed directly into the
Brain density field as S_stress(x, y).

Usage:
    from optimization_engine.extractors.extract_stress_field_2d import extract_stress_field_2d

    field = extract_stress_field_2d(
        op2_file="path/to/results.op2",
        bdf_file="path/to/model.bdf",
        transform=geometry["transform"],   # from geometry_sandbox_N.json
    )

    # field["nodes_2d"]  → (N, 2) array of [u, v] sandbox coords
    # field["stress"]    → (N,)   array of von Mises stress in MPa
    # field["max_stress"] → peak stress in MPa

Unit Note: NX Nastran in kg-mm-s outputs stress in kPa → divided by 1000 → MPa.
"""

from __future__ import annotations

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

import numpy as np
from pyNastran.bdf.bdf import BDF
from pyNastran.op2.op2 import OP2


# ---------------------------------------------------------------------------
# 3D → 2D coordinate projection
# ---------------------------------------------------------------------------

def _project_to_2d(
    xyz: np.ndarray,
    transform: Dict[str, Any],
) -> np.ndarray:
    """
    Project 3D points onto the sandbox plane using the geometry transform.

    The transform (from geometry_sandbox_N.json) defines:
      - origin:  3D origin of the sandbox plane
      - x_axis:  direction of sandbox U axis in 3D
      - y_axis:  direction of sandbox V axis in 3D
      - normal:  plate normal (thickness direction — discarded)

    Inverse of import_profile.py's unproject_point_to_3d().

    Args:
        xyz:       (N, 3) array of 3D points
        transform: dict with 'origin', 'x_axis', 'y_axis', 'normal'

    Returns:
        (N, 2) array of [u, v] sandbox coordinates
    """
    origin = np.array(transform["origin"])
    x_axis = np.array(transform["x_axis"])
    y_axis = np.array(transform["y_axis"])

    # Translate to origin
    rel = xyz - origin  # (N, 3)

    # Project onto sandbox axes
    u = rel @ x_axis    # dot product with x_axis
    v = rel @ y_axis    # dot product with y_axis

    return np.column_stack([u, v])


# ---------------------------------------------------------------------------
# Element centroid extraction from BDF
# ---------------------------------------------------------------------------

def _get_element_centroids(bdf: BDF) -> Dict[int, np.ndarray]:
    """
    Compute centroid for every element in the BDF model.

    Returns:
        {element_id: centroid_xyz (3,)}
    """
    node_xyz = {nid: np.array(node.xyz) for nid, node in bdf.nodes.items()}
    centroids = {}

    for eid, elem in bdf.elements.items():
        try:
            nids = elem.node_ids
            pts = np.array([node_xyz[n] for n in nids if n in node_xyz])
            if len(pts) > 0:
                centroids[eid] = pts.mean(axis=0)
        except Exception:
            pass

    return centroids


# ---------------------------------------------------------------------------
# Von Mises stress extraction from OP2
# ---------------------------------------------------------------------------

def _get_all_von_mises(
    model: OP2,
    subcase: int,
    convert_to_mpa: bool,
) -> Dict[int, float]:
    """
    Extract von Mises stress for every element across all solid + shell types.

    Returns:
        {element_id: von_mises_stress}
    """
    SOLID_TYPES = ["ctetra", "chexa", "cpenta", "cpyram"]
    SHELL_TYPES = ["cquad4", "ctria3"]
    ALL_TYPES = SOLID_TYPES + SHELL_TYPES

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

    stress_container = model.op2_results.stress
    elem_stress: Dict[int, float] = {}

    for elem_type in ALL_TYPES:
        attr = f"{elem_type}_stress"
        if not hasattr(stress_container, attr):
            continue

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

        available = list(stress_dict.keys())
        if not available:
            continue

        sc = subcase if subcase in available else available[0]
        stress = stress_dict[sc]

        if not stress.is_von_mises:
            continue

        ncols = stress.data.shape[2]
        # Von Mises column: solid=9, shell=7
        vm_col = 9 if ncols >= 10 else 7 if ncols == 8 else ncols - 1

        itime = 0
        von_mises = stress.data[itime, :, vm_col]   # (n_elements,)

        # element_node: list of (eid, node_id) pairs — may repeat for each node
        for i, (eid, _node) in enumerate(stress.element_node):
            vm = float(von_mises[i])
            # Keep max stress if element appears multiple times (e.g. corner nodes)
            if eid not in elem_stress or vm > elem_stress[eid]:
                elem_stress[eid] = vm

    if not elem_stress:
        raise ValueError("No von Mises stress data found in OP2 file")

    if convert_to_mpa:
        elem_stress = {eid: v / 1000.0 for eid, v in elem_stress.items()}

    return elem_stress


# ---------------------------------------------------------------------------
# Main extractor
# ---------------------------------------------------------------------------

def extract_stress_field_2d(
    op2_file: Path,
    bdf_file: Path,
    transform: Dict[str, Any],
    subcase: int = 1,
    convert_to_mpa: bool = True,
    sigma_yield: Optional[float] = None,
) -> Dict[str, Any]:
    """
    Extract a 2D von Mises stress field projected onto the sandbox plane.

    For 3D solid meshes: element centroids are projected to 2D, then stress
    values at the same (u, v) location are averaged through thickness.

    For 2D shell meshes: centroids are directly in-plane, no averaging needed.

    Args:
        op2_file:      Path to NX Nastran OP2 results file
        bdf_file:      Path to BDF model file (for geometry/node positions)
        transform:     Sandbox plane transform dict from geometry_sandbox_N.json
                       Keys: 'origin', 'x_axis', 'y_axis', 'normal'
        subcase:       Subcase ID to extract (default: 1)
        convert_to_mpa: Divide by 1000 to convert NX kPa → MPa (default: True)
        sigma_yield:   Optional yield strength in MPa. If provided, adds a
                       'stress_normalized' field (0..1 scale) for density feedback.

    Returns:
        dict with:
            'nodes_2d':          (N, 2) ndarray — [u, v] in sandbox 2D coords (mm)
            'stress':            (N,)   ndarray — von Mises stress (MPa or kPa)
            'max_stress':        float  — peak stress value
            'mean_stress':       float  — mean stress value
            'percentile_95':     float  — 95th percentile (robust peak)
            'units':             str    — 'MPa' or 'kPa'
            'n_elements':        int    — number of elements with stress data
            'stress_normalized': (N,)   ndarray — stress / sigma_yield (if provided)
            'sigma_yield':       float  — yield strength used (if provided)
    """
    op2_file = Path(op2_file)
    bdf_file = Path(bdf_file)

    # --- Load BDF geometry ---
    bdf = BDF(debug=False)
    bdf.read_bdf(str(bdf_file), xref=True)

    centroids_3d = _get_element_centroids(bdf)  # {eid: xyz}

    # --- Load OP2 stress ---
    model = OP2(debug=False, log=None)
    model.read_op2(str(op2_file))

    elem_stress = _get_all_von_mises(model, subcase, convert_to_mpa)

    # --- Match elements: keep only those with both centroid and stress ---
    common_ids = sorted(set(centroids_3d.keys()) & set(elem_stress.keys()))
    if not common_ids:
        raise ValueError(
            f"No matching elements between BDF ({len(centroids_3d)} elements) "
            f"and OP2 ({len(elem_stress)} elements). Check that they are from the same model."
        )

    xyz_arr    = np.array([centroids_3d[eid] for eid in common_ids])  # (N, 3)
    stress_arr = np.array([elem_stress[eid]  for eid in common_ids])  # (N,)

    # --- Project 3D centroids → 2D sandbox coords ---
    nodes_2d = _project_to_2d(xyz_arr, transform)  # (N, 2)

    # --- For 3D solid meshes: average through-thickness duplicates ---
    # Elements at the same (u, v) xy-location but different thickness positions
    # get averaged to produce a single 2D stress value per location.
    uv_rounded = np.round(nodes_2d, decimals=1)  # group within 0.1mm
    uv_tuples = [tuple(r) for r in uv_rounded]

    unique_uvs: Dict[tuple, list] = {}
    for i, uv in enumerate(uv_tuples):
        unique_uvs.setdefault(uv, []).append(stress_arr[i])

    uv_final    = np.array([list(k) for k in unique_uvs.keys()])
    stress_final = np.array([np.mean(v) for v in unique_uvs.values()])

    n_raw      = len(stress_arr)
    n_averaged = len(stress_final)
    n_layers   = round(n_raw / n_averaged) if n_averaged > 0 else 1

    result = {
        "nodes_2d":      uv_final,
        "stress":        stress_final,
        "max_stress":    float(np.max(stress_final)),
        "mean_stress":   float(np.mean(stress_final)),
        "percentile_95": float(np.percentile(stress_final, 95)),
        "units":         "MPa" if convert_to_mpa else "kPa",
        "n_elements":    n_averaged,
        "n_raw_elements": n_raw,
        "n_thickness_layers": n_layers,
    }

    if sigma_yield is not None:
        result["stress_normalized"] = stress_final / sigma_yield
        result["sigma_yield"] = sigma_yield

    return result


# ---------------------------------------------------------------------------
# Save / load helpers
# ---------------------------------------------------------------------------

def save_stress_field(field: Dict[str, Any], output_path: Path) -> None:
    """
    Save extracted stress field to an NPZ file for fast reloading.

    Usage:
        save_stress_field(field, "trial_0001/stress_field_2d.npz")
    """
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    np.savez(
        str(output_path),
        nodes_2d=field["nodes_2d"],
        stress=field["stress"],
        stress_normalized=field.get("stress_normalized", np.array([])),
        max_stress=field["max_stress"],
        mean_stress=field["mean_stress"],
        percentile_95=field["percentile_95"],
        sigma_yield=field.get("sigma_yield", 0.0),
        n_elements=field["n_elements"],
    )


def load_stress_field(npz_path: Path) -> Dict[str, Any]:
    """
    Load a previously saved stress field.

    Usage:
        field = load_stress_field("trial_0001/stress_field_2d.npz")
    """
    data = np.load(str(npz_path), allow_pickle=False)
    field = {
        "nodes_2d":      data["nodes_2d"],
        "stress":        data["stress"],
        "max_stress":    float(data["max_stress"]),
        "mean_stress":   float(data["mean_stress"]),
        "percentile_95": float(data["percentile_95"]),
        "n_elements":    int(data["n_elements"]),
    }
    if data["sigma_yield"] > 0:
        field["stress_normalized"] = data["stress_normalized"]
        field["sigma_yield"] = float(data["sigma_yield"])
    return field


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import sys
    import json

    if len(sys.argv) < 4:
        print("Usage: python extract_stress_field_2d.py <op2> <bdf> <geometry_sandbox.json> [sigma_yield]")
        sys.exit(1)

    op2  = Path(sys.argv[1])
    bdf  = Path(sys.argv[2])
    geom = Path(sys.argv[3])
    sy   = float(sys.argv[4]) if len(sys.argv) > 4 else None

    with open(geom) as f:
        geometry = json.load(f)

    field = extract_stress_field_2d(op2, bdf, geometry["transform"], sigma_yield=sy)

    print(f"Extracted {field['n_elements']} elements "
          f"(from {field['n_raw_elements']} raw, {field['n_thickness_layers']} thickness layers)")
    print(f"Max stress:    {field['max_stress']:.1f} {field['units']}")
    print(f"Mean stress:   {field['mean_stress']:.1f} {field['units']}")
    print(f"95th pct:      {field['percentile_95']:.1f} {field['units']}")

    out = op2.with_suffix(".stress_field_2d.npz")
    save_stress_field(field, out)
    print(f"Saved to: {out}")