Atomizer/optimization_engine/insights/modal_analysis.py

"""
Modal Analysis Insight

Provides visualization of natural frequencies and mode shapes from FEA results.
Shows animated mode shapes, frequency spectrum, and modal participation factors.

Applicable to: Dynamic/vibration optimization studies.
"""

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 ModalInsight(StudyInsight):
    """
    Modal analysis visualization.

    Shows:
    - Natural frequency spectrum (bar chart)
    - Mode shape visualization (3D deformed mesh)
    - Mode description table
    - Frequency vs mode number plot
    """

    insight_type = "modal"
    name = "Modal Analysis"
    description = "Natural frequencies and mode shapes visualization"
    category = "structural_modal"
    applicable_to = ["modal", "vibration", "dynamic", "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._eigenvectors: Optional[Dict] = None
        self._frequencies: Optional[List] = None

    def can_generate(self) -> bool:
        """Check if OP2 file with eigenvalue/eigenvector 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 modal data exists
        try:
            _load_dependencies()
            op2 = _OP2()
            op2.read_op2(str(self.op2_path))
            return bool(op2.eigenvectors)
        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 modal data from OP2."""
        if self._eigenvectors 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 modal data
        op2 = _OP2()
        op2.read_op2(str(self.op2_path))

        self._eigenvectors = {}
        self._frequencies = []

        for key, eig in op2.eigenvectors.items():
            # Get frequencies
            if hasattr(eig, 'modes') and hasattr(eig, 'cycles'):
                modes = eig.modes
                freqs = eig.cycles  # Frequencies in Hz
            elif hasattr(eig, 'eigrs'):
                # Eigenvalues (radians/sec)^2
                eigrs = eig.eigrs
                freqs = np.sqrt(np.abs(eigrs)) / (2 * np.pi)
                modes = list(range(1, len(freqs) + 1))
            else:
                continue

            for i, (mode, freq) in enumerate(zip(modes, freqs)):
                self._frequencies.append({
                    'mode': int(mode),
                    'frequency_hz': float(freq),
                })

            # Get mode shapes
            if hasattr(eig, 'data'):
                data = eig.data
                ngt = eig.node_gridtype.astype(int)
                node_ids = ngt if ngt.ndim == 1 else ngt[:, 0]

                for mode_idx, mode_num in enumerate(modes):
                    if data.ndim == 3:
                        mode_data = data[mode_idx]
                    else:
                        mode_data = data

                    self._eigenvectors[int(mode_num)] = {
                        'node_ids': node_ids,
                        'displacements': mode_data.copy(),
                    }

        # Sort frequencies by mode number
        self._frequencies.sort(key=lambda x: x['mode'])

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

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

        _load_dependencies()

        # Configuration
        n_modes_show = config.extra.get('n_modes', 20)
        mode_to_show = config.extra.get('show_mode', 1)  # Which mode shape to display
        deform_scale = config.amplification if config.amplification != 1.0 else 50.0

        # Limit to available modes
        freq_data = self._frequencies[:n_modes_show]
        modes = [f['mode'] for f in freq_data]
        frequencies = [f['frequency_hz'] for f in freq_data]

        # Build visualization
        fig = _make_subplots(
            rows=2, cols=2,
            specs=[
                [{"type": "scene"}, {"type": "xy"}],
                [{"type": "xy"}, {"type": "table"}]
            ],
            subplot_titles=[
                f"<b>Mode {mode_to_show} Shape</b>",
                "<b>Natural Frequencies</b>",
                "<b>Frequency Spectrum</b>",
                "<b>Mode Summary</b>"
            ]
        )

        # Mode shape (3D)
        if self._node_geo and mode_to_show in self._eigenvectors:
            mode_data = self._eigenvectors[mode_to_show]
            node_ids = mode_data['node_ids']
            disps = mode_data['displacements']

            X, Y, Z = [], [], []
            Xd, Yd, Zd = [], [], []
            colors = []

            for nid, disp in zip(node_ids, disps):
                geo = self._node_geo.get(int(nid))
                if geo is None:
                    continue

                X.append(geo[0])
                Y.append(geo[1])
                Z.append(geo[2])

                # Deformed position
                Xd.append(geo[0] + deform_scale * disp[0])
                Yd.append(geo[1] + deform_scale * disp[1])
                Zd.append(geo[2] + deform_scale * disp[2])

                # Color by displacement magnitude
                mag = np.sqrt(disp[0]**2 + disp[1]**2 + disp[2]**2)
                colors.append(mag)

            X, Y, Z = np.array(X), np.array(Y), np.array(Z)
            Xd, Yd, Zd = np.array(Xd), np.array(Yd), np.array(Zd)
            colors = np.array(colors)

            # Try to create mesh
            try:
                tri = _Triangulation(Xd, Yd)
                if tri.triangles is not None and len(tri.triangles) > 0:
                    i, j, k = tri.triangles.T
                    fig.add_trace(_go.Mesh3d(
                        x=Xd, y=Yd, z=Zd,
                        i=i, j=j, k=k,
                        intensity=colors,
                        colorscale='Viridis',
                        opacity=0.9,
                        flatshading=False,
                        showscale=True,
                        colorbar=dict(title="Disp. Mag.", thickness=10, len=0.4)
                    ), row=1, col=1)
            except Exception:
                # Fallback: scatter
                fig.add_trace(_go.Scatter3d(
                    x=Xd, y=Yd, z=Zd,
                    mode='markers',
                    marker=dict(size=3, color=colors, colorscale='Viridis', showscale=True),
                ), row=1, col=1)

        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),
        )

        # Frequency bar chart
        fig.add_trace(_go.Bar(
            x=modes,
            y=frequencies,
            marker_color='#3b82f6',
            text=[f"{f:.1f} Hz" for f in frequencies],
            textposition='outside',
            name='Frequency'
        ), row=1, col=2)

        fig.update_xaxes(title_text="Mode Number", row=1, col=2)
        fig.update_yaxes(title_text="Frequency (Hz)", row=1, col=2)

        # Frequency spectrum (log scale)
        fig.add_trace(_go.Scatter(
            x=modes,
            y=frequencies,
            mode='lines+markers',
            marker=dict(size=8, color='#22c55e'),
            line=dict(width=2, color='#22c55e'),
            name='Frequency'
        ), row=2, col=1)

        fig.update_xaxes(title_text="Mode Number", row=2, col=1)
        fig.update_yaxes(title_text="Frequency (Hz)", type='log', row=2, col=1)

        # Summary table
        mode_labels = [f"Mode {m}" for m in modes[:10]]
        freq_labels = [f"{f:.2f} Hz" for f in frequencies[:10]]

        fig.add_trace(_go.Table(
            header=dict(values=["<b>Mode</b>", "<b>Frequency</b>"],
                        fill_color='#1f2937', font=dict(color='white')),
            cells=dict(values=[mode_labels, freq_labels],
                       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 Modal 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"modal_{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={
                'n_modes': len(self._frequencies),
                'frequencies_hz': frequencies,
                'first_frequency_hz': frequencies[0] if frequencies else None,
                'shown_mode': mode_to_show,
            }
        )