Atomizer/optimization_engine/insights/base.py

"""
Study Insights - Base Classes and Infrastructure

Study Insights provide physics-focused visualizations for optimization results.
Unlike Analysis (optimizer-centric), Insights show the engineering reality
of specific designs.

Architecture:
- StudyInsight: Abstract base class for all insight types
- InsightRegistry: Central registry for available insight types
- Each insight can generate standalone HTML or Plotly data for dashboard

Usage:
    from optimization_engine.insights import get_insight, list_insights

    # Get specific insight
    insight = get_insight('zernike_wfe')
    if insight.can_generate(study_path):
        html_path = insight.generate_html(study_path, trial_id=47)
        plotly_data = insight.get_plotly_data(study_path, trial_id=47)
"""

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional, Type
import json


@dataclass
class InsightConfig:
    """Configuration for an insight instance."""
    trial_id: Optional[int] = None  # Specific trial to visualize (None = best)
    colorscale: str = 'Turbo'
    output_dir: Optional[Path] = None  # Where to save HTML (None = study/3_insights/)

    # Visual settings
    amplification: float = 1.0  # Deformation scale factor
    lighting: bool = True  # 3D lighting effects

    # Type-specific config (passed through)
    extra: Dict[str, Any] = field(default_factory=dict)


@dataclass
class InsightResult:
    """Result from generating an insight."""
    success: bool
    html_path: Optional[Path] = None
    plotly_figure: Optional[Dict[str, Any]] = None  # Plotly figure as dict
    summary: Optional[Dict[str, Any]] = None  # Key metrics
    error: Optional[str] = None


class StudyInsight(ABC):
    """
    Abstract base class for study-specific physics visualizations.

    Each insight type provides:
    - Detection: Can this insight be generated for this study?
    - HTML generation: Standalone interactive report
    - Plotly data: For embedding in dashboard
    - Summary: Key metrics extracted

    Subclasses must implement:
    - insight_type: Unique identifier (e.g., 'zernike_wfe')
    - name: Human-readable name
    - description: What this insight shows
    - applicable_to: List of study types this applies to
    - can_generate(): Check if study has required data
    - _generate(): Core generation logic
    """

    # Class-level metadata (override in subclasses)
    insight_type: str = "base"
    name: str = "Base Insight"
    description: str = "Abstract base insight"
    applicable_to: List[str] = []  # e.g., ['mirror', 'structural', 'all']

    # Required files/data patterns
    required_files: List[str] = []  # e.g., ['*.op2', '*.bdf']

    def __init__(self, study_path: Path):
        """
        Initialize insight for a specific study.

        Args:
            study_path: Path to study directory (studies/{name}/)
        """
        self.study_path = Path(study_path)
        self.setup_path = self.study_path / "1_setup"
        self.results_path = self.study_path / "2_results"
        self.insights_path = self.study_path / "3_insights"

        # Load study config if available
        self.config = self._load_study_config()

    def _load_study_config(self) -> Dict[str, Any]:
        """Load optimization_config.json if it exists."""
        config_path = self.setup_path / "optimization_config.json"
        if config_path.exists():
            with open(config_path) as f:
                return json.load(f)
        return {}

    @abstractmethod
    def can_generate(self) -> bool:
        """
        Check if this insight can be generated for the study.

        Returns:
            True if all required data is available
        """
        pass

    @abstractmethod
    def _generate(self, config: InsightConfig) -> InsightResult:
        """
        Core generation logic. Implemented by subclasses.

        Args:
            config: Insight configuration

        Returns:
            InsightResult with HTML path and/or Plotly data
        """
        pass

    def generate(self, config: Optional[InsightConfig] = None) -> InsightResult:
        """
        Generate the insight visualization.

        Args:
            config: Optional configuration (uses defaults if None)

        Returns:
            InsightResult with generated content
        """
        if config is None:
            config = InsightConfig()

        # Ensure output directory exists
        if config.output_dir is None:
            config.output_dir = self.insights_path
        config.output_dir.mkdir(parents=True, exist_ok=True)

        # Check prerequisites
        if not self.can_generate():
            return InsightResult(
                success=False,
                error=f"Cannot generate {self.name}: required data not found"
            )

        try:
            return self._generate(config)
        except Exception as e:
            return InsightResult(
                success=False,
                error=f"Error generating {self.name}: {str(e)}"
            )

    def generate_html(
        self,
        trial_id: Optional[int] = None,
        **kwargs
    ) -> Optional[Path]:
        """
        Convenience method to generate standalone HTML.

        Args:
            trial_id: Specific trial to visualize (None = best)
            **kwargs: Additional config options

        Returns:
            Path to generated HTML file, or None on failure
        """
        config = InsightConfig(trial_id=trial_id, extra=kwargs)
        result = self.generate(config)
        return result.html_path if result.success else None

    def get_plotly_data(
        self,
        trial_id: Optional[int] = None,
        **kwargs
    ) -> Optional[Dict[str, Any]]:
        """
        Get Plotly figure data for dashboard embedding.

        Args:
            trial_id: Specific trial to visualize (None = best)
            **kwargs: Additional config options

        Returns:
            Plotly figure as dictionary, or None on failure
        """
        config = InsightConfig(trial_id=trial_id, extra=kwargs)
        result = self.generate(config)
        return result.plotly_figure if result.success else None

    def get_summary(self, trial_id: Optional[int] = None) -> Optional[Dict[str, Any]]:
        """
        Get key metrics summary without full visualization.

        Args:
            trial_id: Specific trial (None = best)

        Returns:
            Dictionary of key metrics
        """
        config = InsightConfig(trial_id=trial_id)
        result = self.generate(config)
        return result.summary if result.success else None


class InsightRegistry:
    """
    Central registry for available insight types.

    Usage:
        registry = InsightRegistry()
        registry.register(ZernikeWFEInsight)

        # Get insight for a study
        insight = registry.get('zernike_wfe', study_path)

        # List available insights for a study
        available = registry.list_available(study_path)
    """

    _instance = None
    _insights: Dict[str, Type[StudyInsight]] = {}

    def __new__(cls):
        """Singleton pattern."""
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._insights = {}
        return cls._instance

    def register(self, insight_class: Type[StudyInsight]) -> None:
        """
        Register an insight type.

        Args:
            insight_class: StudyInsight subclass to register
        """
        self._insights[insight_class.insight_type] = insight_class

    def get(self, insight_type: str, study_path: Path) -> Optional[StudyInsight]:
        """
        Get an insight instance for a study.

        Args:
            insight_type: Registered insight type ID
            study_path: Path to study directory

        Returns:
            Configured insight instance, or None if not found
        """
        if insight_type not in self._insights:
            return None
        return self._insights[insight_type](study_path)

    def list_all(self) -> List[Dict[str, Any]]:
        """
        List all registered insight types.

        Returns:
            List of insight metadata dictionaries
        """
        return [
            {
                'type': cls.insight_type,
                'name': cls.name,
                'description': cls.description,
                'applicable_to': cls.applicable_to
            }
            for cls in self._insights.values()
        ]

    def list_available(self, study_path: Path) -> List[Dict[str, Any]]:
        """
        List insights that can be generated for a specific study.

        Args:
            study_path: Path to study directory

        Returns:
            List of available insight metadata
        """
        available = []
        for insight_type, cls in self._insights.items():
            try:
                insight = cls(study_path)
                if insight.can_generate():
                    available.append({
                        'type': insight_type,
                        'name': cls.name,
                        'description': cls.description
                    })
            except Exception:
                pass  # Skip insights that fail to initialize
        return available


# Global registry instance
_registry = InsightRegistry()


def register_insight(insight_class: Type[StudyInsight]) -> Type[StudyInsight]:
    """
    Decorator to register an insight class.

    Usage:
        @register_insight
        class MyInsight(StudyInsight):
            insight_type = 'my_insight'
            ...
    """
    _registry.register(insight_class)
    return insight_class


def get_insight(insight_type: str, study_path: Path) -> Optional[StudyInsight]:
    """Get an insight instance by type."""
    return _registry.get(insight_type, study_path)


def list_insights() -> List[Dict[str, Any]]:
    """List all registered insight types."""
    return _registry.list_all()


def list_available_insights(study_path: Path) -> List[Dict[str, Any]]:
    """List insights available for a specific study."""
    return _registry.list_available(study_path)