Atomizer/optimization_engine/interview/interview_intelligence.py

"""
Interview Intelligence

Smart features for the interview process:
- ExtractorMapper: Maps goals to appropriate extractors
- InterviewIntelligence: Auto-detection, inference, complexity determination
"""

from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Any, Optional, Literal, Tuple
import re


@dataclass
class ExtractorSelection:
    """Result of mapping a goal to an extractor."""
    extractor_id: str
    extractor_name: str
    goal_type: str  # minimize, maximize, target
    params: Dict[str, Any] = field(default_factory=dict)
    fallback: Optional[str] = None
    confidence: float = 1.0
    notes: Optional[str] = None


class ExtractorMapper:
    """
    Maps physics goals to appropriate extractors.

    Uses the Atomizer extractor library (SYS_12) to select
    the right extractor for each objective or constraint.
    """

    # Goal to extractor mapping
    GOAL_MAP = {
        # Mass objectives
        "minimize_mass": ExtractorSelection(
            extractor_id="E4",
            extractor_name="BDF Mass Extraction",
            goal_type="minimize",
            fallback="E5",
            notes="Uses BDF parsing for accurate mass. Falls back to NX expression."
        ),
        "minimize_weight": ExtractorSelection(
            extractor_id="E4",
            extractor_name="BDF Mass Extraction",
            goal_type="minimize",
            fallback="E5"
        ),

        # Displacement/stiffness objectives
        "minimize_displacement": ExtractorSelection(
            extractor_id="E1",
            extractor_name="Displacement Extraction",
            goal_type="minimize",
            params={"component": "magnitude", "node_id": "auto"},
            notes="Extracts displacement magnitude. Node ID auto-detected from max."
        ),
        "maximize_stiffness": ExtractorSelection(
            extractor_id="E1",
            extractor_name="Displacement Extraction",
            goal_type="minimize",  # Stiffness = 1/displacement
            params={"component": "magnitude", "node_id": "auto"},
            notes="Stiffness maximization = displacement minimization"
        ),

        # Frequency objectives
        "maximize_frequency": ExtractorSelection(
            extractor_id="E2",
            extractor_name="Frequency Extraction",
            goal_type="maximize",
            params={"mode_number": 1},
            notes="First natural frequency. Mode number adjustable."
        ),
        "target_frequency": ExtractorSelection(
            extractor_id="E2",
            extractor_name="Frequency Extraction",
            goal_type="target",
            params={"mode_number": 1, "target": None},
            notes="Target a specific frequency value."
        ),

        # Stress objectives
        "minimize_stress": ExtractorSelection(
            extractor_id="E3",
            extractor_name="Solid Stress Extraction",
            goal_type="minimize",
            params={"element_type": "auto", "stress_type": "von_mises"},
            notes="Von Mises stress. Element type auto-detected."
        ),

        # Optical objectives
        "minimize_wavefront_error": ExtractorSelection(
            extractor_id="E8",
            extractor_name="Zernike Wavefront Fitting",
            goal_type="minimize",
            params={"n_terms": 15, "radius": "auto"},
            notes="Fits surface to Zernike polynomials. Optical applications."
        ),

        # Custom
        "custom": ExtractorSelection(
            extractor_id="custom",
            extractor_name="Custom Extractor",
            goal_type="custom",
            confidence=0.5,
            notes="User will define custom extraction logic."
        ),
    }

    # Constraint type to extractor mapping
    CONSTRAINT_MAP = {
        "stress": ExtractorSelection(
            extractor_id="E3",
            extractor_name="Solid Stress Extraction",
            goal_type="max",
            params={"stress_type": "von_mises"}
        ),
        "max_stress": ExtractorSelection(
            extractor_id="E3",
            extractor_name="Solid Stress Extraction",
            goal_type="max",
            params={"stress_type": "von_mises"}
        ),
        "displacement": ExtractorSelection(
            extractor_id="E1",
            extractor_name="Displacement Extraction",
            goal_type="max",
            params={"component": "magnitude"}
        ),
        "max_displacement": ExtractorSelection(
            extractor_id="E1",
            extractor_name="Displacement Extraction",
            goal_type="max",
            params={"component": "magnitude"}
        ),
        "frequency": ExtractorSelection(
            extractor_id="E2",
            extractor_name="Frequency Extraction",
            goal_type="min",
            params={"mode_number": 1}
        ),
        "min_frequency": ExtractorSelection(
            extractor_id="E2",
            extractor_name="Frequency Extraction",
            goal_type="min",
            params={"mode_number": 1}
        ),
        "mass": ExtractorSelection(
            extractor_id="E4",
            extractor_name="BDF Mass Extraction",
            goal_type="max"
        ),
        "max_mass": ExtractorSelection(
            extractor_id="E4",
            extractor_name="BDF Mass Extraction",
            goal_type="max"
        ),
    }

    def map_goal_to_extractor(
        self,
        goal: str,
        introspection: Optional[Dict[str, Any]] = None
    ) -> ExtractorSelection:
        """
        Map a physics goal to the appropriate extractor.

        Args:
            goal: Goal identifier (e.g., "minimize_mass")
            introspection: Optional introspection results for auto-detection

        Returns:
            ExtractorSelection with extractor details
        """
        goal_lower = goal.lower().strip()

        # Direct match
        if goal_lower in self.GOAL_MAP:
            selection = self.GOAL_MAP[goal_lower]

            # Auto-detect parameters if introspection available
            if introspection:
                selection = self._refine_selection(selection, introspection)

            return selection

        # Fuzzy matching for common variations
        for key, selection in self.GOAL_MAP.items():
            if key.replace("_", " ") in goal_lower or goal_lower in key:
                return selection

        # Default to custom
        return self.GOAL_MAP["custom"]

    def map_constraint_to_extractor(
        self,
        constraint_type: str,
        introspection: Optional[Dict[str, Any]] = None
    ) -> ExtractorSelection:
        """
        Map a constraint type to the appropriate extractor.

        Args:
            constraint_type: Constraint type (e.g., "stress", "displacement")
            introspection: Optional introspection results

        Returns:
            ExtractorSelection with extractor details
        """
        type_lower = constraint_type.lower().strip()

        if type_lower in self.CONSTRAINT_MAP:
            selection = self.CONSTRAINT_MAP[type_lower]

            if introspection:
                selection = self._refine_selection(selection, introspection)

            return selection

        # Try to infer from name
        if "stress" in type_lower:
            return self.CONSTRAINT_MAP["stress"]
        if "disp" in type_lower or "deflect" in type_lower:
            return self.CONSTRAINT_MAP["displacement"]
        if "freq" in type_lower or "modal" in type_lower:
            return self.CONSTRAINT_MAP["frequency"]
        if "mass" in type_lower or "weight" in type_lower:
            return self.CONSTRAINT_MAP["mass"]

        return ExtractorSelection(
            extractor_id="custom",
            extractor_name="Custom Constraint",
            goal_type="constraint",
            confidence=0.5
        )

    def _refine_selection(
        self,
        selection: ExtractorSelection,
        introspection: Dict[str, Any]
    ) -> ExtractorSelection:
        """Refine extractor selection based on introspection."""
        import copy
        refined = copy.deepcopy(selection)

        # Auto-detect element type for stress extraction
        if refined.extractor_id == "E3" and refined.params.get("element_type") == "auto":
            element_types = introspection.get("element_types", [])
            if "solid" in element_types or any("TET" in e or "HEX" in e for e in element_types):
                refined.params["element_type"] = "solid"
            elif "shell" in element_types or any("QUAD" in e or "TRI" in e for e in element_types):
                refined.params["element_type"] = "shell"
                refined.extractor_id = "E3_shell"  # Use shell stress extractor

        # Auto-detect node for displacement
        if refined.extractor_id == "E1" and refined.params.get("node_id") == "auto":
            # Use max displacement node from baseline if available
            if "max_disp_node" in introspection:
                refined.params["node_id"] = introspection["max_disp_node"]

        return refined

    def get_extractor_summary(self, selections: List[ExtractorSelection]) -> str:
        """Generate a summary of selected extractors."""
        lines = ["**Selected Extractors:**", ""]

        for sel in selections:
            params_str = ""
            if sel.params:
                params_str = " (" + ", ".join(f"{k}={v}" for k, v in sel.params.items()) + ")"

            lines.append(f"- **{sel.extractor_id}**: {sel.extractor_name}{params_str}")
            if sel.notes:
                lines.append(f"  > {sel.notes}")

        return "\n".join(lines)


@dataclass
class StudyTypeInference:
    """Result of inferring study type."""
    study_type: str  # single_objective, multi_objective, parametric
    protocol: str  # protocol_10_single, protocol_11_multi
    confidence: float
    reasons: List[str] = field(default_factory=list)


class InterviewIntelligence:
    """
    Smart features for the interview process.

    Provides:
    - Study type inference from context
    - Auto-selection of extractors
    - History-based suggestions
    - Complexity determination
    """

    def __init__(self):
        """Initialize intelligence module."""
        self.extractor_mapper = ExtractorMapper()

    def infer_study_type(
        self,
        study_name: str,
        user_description: str,
        introspection: Optional[Dict[str, Any]] = None
    ) -> StudyTypeInference:
        """
        Infer study type from available context.

        Args:
            study_name: Study name (may contain hints)
            user_description: User's problem description
            introspection: Optional introspection results

        Returns:
            StudyTypeInference with type and protocol
        """
        reasons = []
        score_multi = 0
        score_single = 0

        text = f"{study_name} {user_description}".lower()

        # Check for multi-objective keywords
        if any(kw in text for kw in ["pareto", "trade-off", "tradeoff", "multi-objective", "multiobjective"]):
            score_multi += 2
            reasons.append("Multi-objective keywords detected")

        if any(kw in text for kw in ["versus", " vs ", "and minimize", "and maximize", "balance"]):
            score_multi += 1
            reasons.append("Conflicting goals language detected")

        # Check for single-objective keywords
        if any(kw in text for kw in ["minimize", "maximize", "reduce", "increase"]):
            # Count occurrences
            count = sum(1 for kw in ["minimize", "maximize", "reduce", "increase"] if kw in text)
            if count == 1:
                score_single += 1
                reasons.append("Single optimization goal language")
            else:
                score_multi += 1
                reasons.append("Multiple optimization verbs detected")

        # Default to single objective if no strong signals
        if score_multi > score_single:
            return StudyTypeInference(
                study_type="multi_objective",
                protocol="protocol_11_multi",
                confidence=min(1.0, 0.5 + score_multi * 0.2),
                reasons=reasons
            )
        else:
            return StudyTypeInference(
                study_type="single_objective",
                protocol="protocol_10_single",
                confidence=min(1.0, 0.5 + score_single * 0.2),
                reasons=reasons if reasons else ["Default to single-objective"]
            )

    def auto_select_extractors(
        self,
        objectives: List[Dict[str, Any]],
        constraints: List[Dict[str, Any]],
        introspection: Optional[Dict[str, Any]] = None
    ) -> Dict[str, ExtractorSelection]:
        """
        Automatically select appropriate extractors.

        Args:
            objectives: List of objective definitions
            constraints: List of constraint definitions
            introspection: Optional introspection results

        Returns:
            Dict mapping objective/constraint names to ExtractorSelection
        """
        selections = {}

        # Map objectives
        for i, obj in enumerate(objectives):
            goal = obj.get("goal", "") if isinstance(obj, dict) else str(obj)
            name = obj.get("name", f"objective_{i}") if isinstance(obj, dict) else f"objective_{i}"

            selection = self.extractor_mapper.map_goal_to_extractor(goal, introspection)
            selections[name] = selection

        # Map constraints
        for i, con in enumerate(constraints):
            con_type = con.get("type", "") if isinstance(con, dict) else str(con)
            name = con.get("name", f"constraint_{i}") if isinstance(con, dict) else f"constraint_{i}"

            selection = self.extractor_mapper.map_constraint_to_extractor(con_type, introspection)
            selections[name] = selection

        return selections

    def determine_complexity(
        self,
        state: "InterviewState",
        introspection: Optional[Dict[str, Any]] = None
    ) -> Literal["simple", "moderate", "complex"]:
        """
        Determine study complexity for adaptive questioning.

        Based on:
        - Number of objectives
        - Number of design variables
        - Analysis complexity
        - Custom components

        Args:
            state: Current interview state
            introspection: Optional introspection results

        Returns:
            Complexity level
        """
        score = 0
        answers = state.answers

        # Objectives
        n_obj = len(answers.get("objectives", []))
        secondary = answers.get("objectives_secondary", [])
        if "none" not in secondary:
            n_obj += len(secondary)

        if n_obj == 1:
            score += 0
        elif n_obj == 2:
            score += 1
        else:
            score += 2

        # Design variables
        n_dvs = len(answers.get("design_variables", []))
        if n_dvs <= 3:
            score += 0
        elif n_dvs <= 6:
            score += 1
        else:
            score += 2

        # Analysis types
        analysis_types = answers.get("analysis_types", [])
        if len(analysis_types) > 2:
            score += 2
        elif len(analysis_types) > 1:
            score += 1

        if "coupled_thermal_structural" in analysis_types:
            score += 1
        if "nonlinear" in analysis_types:
            score += 1

        # Introspection complexity
        if introspection:
            if introspection.get("multiple_solutions", False):
                score += 1
            if len(introspection.get("expressions", [])) > 20:
                score += 1

        # Categorize
        if score <= 2:
            return "simple"
        elif score <= 5:
            return "moderate"
        else:
            return "complex"

    def suggest_trial_count(
        self,
        n_design_variables: int,
        n_objectives: int,
        complexity: str
    ) -> int:
        """
        Suggest appropriate number of trials.

        Args:
            n_design_variables: Number of design variables
            n_objectives: Number of objectives
            complexity: Study complexity level

        Returns:
            Suggested trial count
        """
        # Base: 15 trials per design variable
        base = n_design_variables * 15

        # Multi-objective needs more
        if n_objectives > 1:
            base = int(base * 1.5)

        # Adjust for complexity
        if complexity == "simple":
            base = max(50, base)
        elif complexity == "moderate":
            base = max(100, base)
        else:
            base = max(150, base)

        # Round to nice numbers
        if base <= 50:
            return 50
        elif base <= 75:
            return 75
        elif base <= 100:
            return 100
        elif base <= 150:
            return 150
        elif base <= 200:
            return 200
        else:
            return int((base // 100) * 100)

    def suggest_sampler(
        self,
        n_objectives: int,
        n_design_variables: int
    ) -> str:
        """
        Suggest appropriate sampler/optimizer.

        Args:
            n_objectives: Number of objectives
            n_design_variables: Number of design variables

        Returns:
            Sampler name
        """
        if n_objectives > 1:
            return "NSGA-II"  # Multi-objective
        elif n_design_variables <= 3:
            return "TPE"  # Tree-structured Parzen Estimator
        elif n_design_variables <= 10:
            return "CMA-ES"  # Covariance Matrix Adaptation
        else:
            return "TPE"  # TPE handles high dimensions well

    def analyze_design_variable_candidates(
        self,
        expressions: List[Dict[str, Any]]
    ) -> List[Dict[str, Any]]:
        """
        Analyze expressions to find design variable candidates.

        Args:
            expressions: List of expressions from introspection

        Returns:
            Sorted list of candidates with scores
        """
        candidates = []

        # High confidence patterns
        high_patterns = [
            (r"thickness", "Thickness parameter"),
            (r"width", "Width parameter"),
            (r"height", "Height parameter"),
            (r"diameter", "Diameter parameter"),
            (r"radius", "Radius parameter"),
            (r"length", "Length parameter"),
            (r"depth", "Depth parameter"),
            (r"angle", "Angle parameter"),
            (r"fillet", "Fillet radius"),
            (r"chamfer", "Chamfer dimension"),
            (r"rib_", "Rib parameter"),
            (r"wall_", "Wall parameter"),
            (r"flange_", "Flange parameter"),
        ]

        # Medium confidence patterns
        medium_patterns = [
            (r"dim_", "Dimension parameter"),
            (r"size_", "Size parameter"),
            (r"param_", "Named parameter"),
            (r"^p\d+$", "Numbered parameter"),
            (r"var_", "Variable"),
        ]

        # Exclusion patterns
        exclude_patterns = [
            r"mesh_", r"count_", r"num_", r"material",
            r"derived_", r"calc_", r"_result$", r"_output$",
            r"^n\d+$", r"count$"
        ]

        for expr in expressions:
            name = expr.get("name", "")
            value = expr.get("value")
            formula = expr.get("formula", "")

            # Skip non-numeric
            if not isinstance(value, (int, float)):
                continue

            # Skip formulas (computed values)
            if formula and formula != str(value):
                continue

            # Check exclusions
            if any(re.search(p, name.lower()) for p in exclude_patterns):
                continue

            # Score
            score = 0
            reason = "Named expression"

            for pattern, desc in high_patterns:
                if re.search(pattern, name.lower()):
                    score = 3
                    reason = desc
                    break

            if score == 0:
                for pattern, desc in medium_patterns:
                    if re.search(pattern, name.lower()):
                        score = 2
                        reason = desc
                        break

            if score == 0 and len(name) > 2:
                score = 1

            if score > 0:
                candidates.append({
                    "name": name,
                    "value": value,
                    "score": score,
                    "reason": reason,
                    "suggested_min": round(value * 0.5, 3) if value > 0 else round(value * 1.5, 3),
                    "suggested_max": round(value * 1.5, 3) if value > 0 else round(value * 0.5, 3),
                })

        # Sort by score descending
        candidates.sort(key=lambda x: (-x["score"], x["name"]))

        return candidates


# Import for type hints
from typing import TYPE_CHECKING
if TYPE_CHECKING:
    from .interview_state import InterviewState