Atomizer/optimization_engine/insights/thermal_field.py

"""
Thermal Field Insight

Provides visualization of temperature distributions from thermal FEA results.
Shows temperature contours, gradients, and thermal statistics.

Applicable to: Thermal analysis and thermo-structural 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 ThermalInsight(StudyInsight):
    """
    Thermal field visualization.

    Shows:
    - 3D mesh colored by temperature
    - Temperature distribution histogram
    - Hot/cold spot identification
    - Temperature gradient visualization
    """

    insight_type = "thermal"
    name = "Thermal Analysis"
    description = "Temperature distribution and thermal gradients"
    category = "thermal"
    applicable_to = ["thermal", "thermo-structural", "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._temperatures: Optional[Dict] = None

    def can_generate(self) -> bool:
        """Check if OP2 file with temperature 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 temperature data exists
        try:
            _load_dependencies()
            op2 = _OP2()
            op2.read_op2(str(self.op2_path))
            # Check for temperature results (various possible attributes)
            return bool(hasattr(op2, 'temperatures') and op2.temperatures)
        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 temperature data from OP2."""
        if self._temperatures 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 temperature data
        op2 = _OP2()
        op2.read_op2(str(self.op2_path))

        self._temperatures = {}

        if hasattr(op2, 'temperatures'):
            for key, temp_obj in op2.temperatures.items():
                if hasattr(temp_obj, 'data'):
                    data = temp_obj.data
                    if data.ndim == 3:
                        data = data[0]  # First time step

                    ngt = temp_obj.node_gridtype.astype(int) if hasattr(temp_obj, 'node_gridtype') else None
                    node_ids = ngt if ngt is not None and ngt.ndim == 1 else (ngt[:, 0] if ngt is not None else None)

                    self._temperatures[str(key)] = {
                        'node_ids': node_ids,
                        'temperatures': data.flatten() if data.ndim > 1 else data,
                    }

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

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

        _load_dependencies()

        # Configuration
        colorscale = config.extra.get('colorscale', 'Thermal')
        temp_unit = config.extra.get('temp_unit', 'K')

        # Aggregate temperature data
        all_temps = []
        all_node_ids = []
        for key, data in self._temperatures.items():
            temps = data['temperatures']
            if isinstance(temps, np.ndarray):
                all_temps.extend(temps.flatten().tolist())
                if data['node_ids'] is not None:
                    all_node_ids.extend(data['node_ids'].flatten().tolist())

        all_temps = np.array(all_temps)
        if len(all_temps) == 0:
            return InsightResult(success=False, error="No valid temperature values found")

        max_temp = float(np.max(all_temps))
        min_temp = float(np.min(all_temps))
        mean_temp = float(np.mean(all_temps))
        temp_range = max_temp - min_temp

        # 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>Temperature Distribution</b>",
                "<b>Temperature Histogram</b>",
                "<b>Summary Statistics</b>"
            ]
        )

        # 3D temperature field
        if self._node_geo and all_node_ids:
            # Build node-to-temp mapping
            temp_map = {}
            for key, data in self._temperatures.items():
                if data['node_ids'] is not None:
                    for nid, temp in zip(data['node_ids'].flatten(), data['temperatures'].flatten()):
                        temp_map[int(nid)] = temp

            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])
            colors = np.array([temp_map.get(nid, mean_temp) for nid in node_ids])

            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.95,
                        flatshading=False,
                        lighting=dict(ambient=0.5, diffuse=0.7, specular=0.2),
                        showscale=True,
                        colorbar=dict(title=f"Temp ({temp_unit})",
                                      thickness=15, len=0.5)
                    ), row=1, col=1)
            except Exception:
                fig.add_trace(_go.Scatter3d(
                    x=X, y=Y, z=Z,
                    mode='markers',
                    marker=dict(size=4, color=colors, colorscale=colorscale, showscale=True),
                ), row=1, col=1)
        else:
            fig.add_annotation(
                text="3D mesh not available",
                xref="paper", yref="paper", x=0.5, y=0.7,
                showarrow=False, font=dict(size=14, color='white')
            )

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

        # Temperature histogram
        fig.add_trace(_go.Histogram(
            x=all_temps,
            nbinsx=50,
            marker_color='#f97316',
            opacity=0.8,
            name='Temperature'
        ), row=2, col=1)

        fig.update_xaxes(title_text=f"Temperature ({temp_unit})", row=2, col=1)
        fig.update_yaxes(title_text="Count", row=2, col=1)

        # Summary table
        stats_labels = [
            "Maximum Temperature",
            "Minimum Temperature",
            "Mean Temperature",
            "Temperature Range",
            "Number of Nodes",
        ]
        stats_values = [
            f"{max_temp:.2f} {temp_unit}",
            f"{min_temp:.2f} {temp_unit}",
            f"{mean_temp:.2f} {temp_unit}",
            f"{temp_range:.2f} {temp_unit}",
            str(len(all_temps)),
        ]

        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 Thermal 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"thermal_{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_temp': max_temp,
                'min_temp': min_temp,
                'mean_temp': mean_temp,
                'temp_range': temp_range,
                'temp_unit': temp_unit,
            }
        )