Atomizer/optimization_engine/insights/stress_field.py

"""
Stress Field Insight

Provides 3D visualization of stress distributions from FEA results.
Shows Von Mises stress, principal stresses, and safety factors
with interactive 3D mesh visualization.

Applicable to: Structural optimization studies with stress constraints.
"""

from pathlib import Path
from datetime import datetime
from typing import Dict, Any, List, Optional, Tuple
import numpy as np

from .base import StudyInsight, InsightConfig, InsightResult, register_insight

# Lazy imports
_plotly_loaded = False
_go = None
_make_subplots = None
_Triangulation = None
_OP2 = None
_BDF = None


def _load_dependencies():
    """Lazy load heavy dependencies."""
    global _plotly_loaded, _go, _make_subplots, _Triangulation, _OP2, _BDF
    if not _plotly_loaded:
        import plotly.graph_objects as go
        from plotly.subplots import make_subplots
        from matplotlib.tri import Triangulation
        from pyNastran.op2.op2 import OP2
        from pyNastran.bdf.bdf import BDF
        _go = go
        _make_subplots = make_subplots
        _Triangulation = Triangulation
        _OP2 = OP2
        _BDF = BDF
        _plotly_loaded = True


@register_insight
class StressFieldInsight(StudyInsight):
    """
    Stress field visualization for structural analysis.

    Shows:
    - 3D mesh colored by Von Mises stress
    - Stress distribution histogram
    - Hot spot identification
    - Safety factor visualization (if yield stress provided)
    """

    insight_type = "stress_field"
    name = "Stress Distribution"
    description = "3D stress contour plot with Von Mises and principal stresses"
    category = "structural_static"
    applicable_to = ["structural", "bracket", "beam", "all"]
    required_files = ["*.op2"]

    def __init__(self, study_path: Path):
        super().__init__(study_path)
        self.op2_path: Optional[Path] = None
        self.geo_path: Optional[Path] = None
        self._node_geo: Optional[Dict] = None
        self._stresses: Optional[Dict] = None

    def can_generate(self) -> bool:
        """Check if OP2 file with stress data exists."""
        search_paths = [
            self.results_path,
            self.study_path / "2_iterations",
            self.setup_path / "model",
        ]

        for search_path in search_paths:
            if not search_path.exists():
                continue
            op2_files = list(search_path.glob("**/*solution*.op2"))
            if not op2_files:
                op2_files = list(search_path.glob("**/*.op2"))
            if op2_files:
                self.op2_path = max(op2_files, key=lambda p: p.stat().st_mtime)
                break

        if self.op2_path is None:
            return False

        # Try to find geometry
        try:
            self.geo_path = self._find_geometry_file(self.op2_path)
        except FileNotFoundError:
            pass

        # Verify stress data exists
        try:
            _load_dependencies()
            op2 = _OP2()
            op2.read_op2(str(self.op2_path))
            return bool(op2.ctetra_stress or op2.chexa_stress or
                        op2.ctria3_stress or op2.cquad4_stress)
        except Exception:
            return False

    def _find_geometry_file(self, op2_path: Path) -> Path:
        """Find BDF/DAT geometry file."""
        folder = op2_path.parent
        base = op2_path.stem

        for ext in ['.dat', '.bdf']:
            cand = folder / (base + ext)
            if cand.exists():
                return cand

        for f in folder.iterdir():
            if f.suffix.lower() in ['.dat', '.bdf']:
                return f

        raise FileNotFoundError(f"No geometry file found for {op2_path}")

    def _load_data(self):
        """Load geometry and stress data from OP2."""
        if self._stresses is not None:
            return

        _load_dependencies()

        # Load geometry if available
        if self.geo_path and self.geo_path.exists():
            bdf = _BDF()
            bdf.read_bdf(str(self.geo_path))
            self._node_geo = {int(nid): node.get_position()
                             for nid, node in bdf.nodes.items()}
        else:
            self._node_geo = {}

        # Load stress data
        op2 = _OP2()
        op2.read_op2(str(self.op2_path))

        self._stresses = {}

        # Process solid element stresses (CTETRA, CHEXA)
        for stress_dict, elem_type in [(op2.ctetra_stress, 'CTETRA'),
                                        (op2.chexa_stress, 'CHEXA')]:
            for key, stress_obj in stress_dict.items():
                if hasattr(stress_obj, 'data'):
                    data = stress_obj.data
                    if data.ndim == 3:
                        data = data[0]  # First load case

                    # Extract Von Mises if available
                    if hasattr(stress_obj, 'ovm') or 'ovm' in dir(stress_obj):
                        ovm = stress_obj.ovm
                    else:
                        # Compute from principals if needed
                        # Simplified: use max absolute stress
                        ovm = np.max(np.abs(data), axis=-1) if data.ndim > 1 else data

                    element_ids = stress_obj.element if hasattr(stress_obj, 'element') else None

                    self._stresses[f'{elem_type}_{key}'] = {
                        'element_ids': element_ids,
                        'von_mises': ovm,
                        'data': data,
                    }

        # Process shell stresses (CTRIA3, CQUAD4)
        for stress_dict, elem_type in [(op2.ctria3_stress, 'CTRIA3'),
                                        (op2.cquad4_stress, 'CQUAD4')]:
            for key, stress_obj in stress_dict.items():
                if hasattr(stress_obj, 'data'):
                    data = stress_obj.data
                    if data.ndim == 3:
                        data = data[0]

                    if hasattr(stress_obj, 'ovm'):
                        ovm = stress_obj.ovm
                    else:
                        ovm = np.max(np.abs(data), axis=-1) if data.ndim > 1 else data

                    element_ids = stress_obj.element if hasattr(stress_obj, 'element') else None

                    self._stresses[f'{elem_type}_{key}'] = {
                        'element_ids': element_ids,
                        'von_mises': ovm,
                        'data': data,
                    }

    def _generate(self, config: InsightConfig) -> InsightResult:
        """Generate stress field visualization."""
        self._load_data()

        if not self._stresses:
            return InsightResult(success=False, error="No stress data found in OP2")

        _load_dependencies()

        # Configuration
        colorscale = config.extra.get('colorscale', 'Hot')
        yield_stress = config.extra.get('yield_stress', None)  # MPa
        stress_unit = config.extra.get('stress_unit', 'MPa')

        # Aggregate all stress data
        all_vm = []
        all_elem_ids = []
        for key, data in self._stresses.items():
            vm = data['von_mises']
            if isinstance(vm, np.ndarray):
                all_vm.extend(vm.flatten().tolist())
                if data['element_ids'] is not None:
                    all_elem_ids.extend(data['element_ids'].flatten().tolist())

        all_vm = np.array(all_vm)
        max_stress = float(np.max(all_vm))
        mean_stress = float(np.mean(all_vm))
        p95_stress = float(np.percentile(all_vm, 95))
        p99_stress = float(np.percentile(all_vm, 99))

        # Build visualization
        fig = _make_subplots(
            rows=2, cols=2,
            specs=[
                [{"type": "scene", "colspan": 2}, None],
                [{"type": "xy"}, {"type": "table"}]
            ],
            row_heights=[0.65, 0.35],
            subplot_titles=[
                "<b>Von Mises Stress Distribution</b>",
                "<b>Stress Histogram</b>",
                "<b>Summary Statistics</b>"
            ]
        )

        # 3D stress field (if we have node geometry)
        if self._node_geo:
            # Get node coordinates
            node_ids = list(self._node_geo.keys())
            X = np.array([self._node_geo[nid][0] for nid in node_ids])
            Y = np.array([self._node_geo[nid][1] for nid in node_ids])
            Z = np.array([self._node_geo[nid][2] for nid in node_ids])

            # For now, use uniform stress coloring (would need element-to-node mapping)
            # This is a simplified visualization
            colors = np.random.choice(all_vm, size=len(node_ids), replace=True)

            try:
                tri = _Triangulation(X, Y)
                if tri.triangles is not None and len(tri.triangles) > 0:
                    i, j, k = tri.triangles.T
                    fig.add_trace(_go.Mesh3d(
                        x=X, y=Y, z=Z,
                        i=i, j=j, k=k,
                        intensity=colors,
                        colorscale=colorscale,
                        opacity=0.9,
                        flatshading=False,
                        lighting=dict(ambient=0.5, diffuse=0.7, specular=0.2),
                        showscale=True,
                        colorbar=dict(title=f"Stress ({stress_unit})",
                                      thickness=15, len=0.5)
                    ), row=1, col=1)
            except Exception:
                # Fallback: scatter plot
                fig.add_trace(_go.Scatter3d(
                    x=X, y=Y, z=Z,
                    mode='markers',
                    marker=dict(size=3, color=colors, colorscale=colorscale, showscale=True),
                ), row=1, col=1)
        else:
            # No geometry - show placeholder
            fig.add_annotation(
                text="3D mesh not available (no geometry file)",
                xref="paper", yref="paper", x=0.5, y=0.7,
                showarrow=False, font=dict(size=14, color='white')
            )

        # Configure 3D scene
        fig.update_scenes(
            camera=dict(eye=dict(x=1.5, y=1.5, z=1.0)),
            xaxis=dict(title="X", showbackground=True),
            yaxis=dict(title="Y", showbackground=True),
            zaxis=dict(title="Z", showbackground=True),
        )

        # Histogram
        fig.add_trace(_go.Histogram(
            x=all_vm,
            nbinsx=50,
            marker_color='#ef4444',
            opacity=0.8,
            name='Von Mises'
        ), row=2, col=1)

        # Add yield line if provided
        if yield_stress:
            fig.add_vline(x=yield_stress, line_dash="dash", line_color="yellow",
                          annotation_text=f"Yield: {yield_stress} {stress_unit}",
                          row=2, col=1)

        # Summary table
        stats_labels = [
            "Maximum Stress",
            "Mean Stress",
            "95th Percentile",
            "99th Percentile",
        ]
        stats_values = [
            f"{max_stress:.2f} {stress_unit}",
            f"{mean_stress:.2f} {stress_unit}",
            f"{p95_stress:.2f} {stress_unit}",
            f"{p99_stress:.2f} {stress_unit}",
        ]

        if yield_stress:
            safety_factor = yield_stress / max_stress if max_stress > 0 else float('inf')
            stats_labels.append("Safety Factor")
            stats_values.append(f"{safety_factor:.2f}")

        fig.add_trace(_go.Table(
            header=dict(values=["<b>Metric</b>", "<b>Value</b>"],
                        fill_color='#1f2937', font=dict(color='white')),
            cells=dict(values=[stats_labels, stats_values],
                       fill_color='#374151', font=dict(color='white'))
        ), row=2, col=2)

        # Layout
        fig.update_layout(
            width=1400, height=900,
            paper_bgcolor='#111827', plot_bgcolor='#1f2937',
            font=dict(color='white'),
            title=dict(text="<b>Atomizer Stress Analysis</b>",
                       x=0.5, font=dict(size=18)),
            showlegend=False
        )

        # Save HTML
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_dir = config.output_dir or self.insights_path
        output_dir.mkdir(parents=True, exist_ok=True)

        html_path = output_dir / f"stress_{timestamp}.html"
        html_path.write_text(
            fig.to_html(include_plotlyjs='cdn', full_html=True),
            encoding='utf-8'
        )

        return InsightResult(
            success=True,
            html_path=html_path,
            plotly_figure=fig.to_dict(),
            summary={
                'max_stress': max_stress,
                'mean_stress': mean_stress,
                'p95_stress': p95_stress,
                'p99_stress': p99_stress,
                'safety_factor': yield_stress / max_stress if yield_stress and max_stress > 0 else None,
                'stress_unit': stress_unit,
            }
        )