pcb-afm-simulation-system/web/backend/app/services/adaptive_loop.py

"""Adaptive Simulation Loop (P3-M5).

Integrates all P3 components into a complete end-to-end adaptive
simulation closed loop:
1. AI plan generation (natural language -> structured plan)
2. L0 pre-screening (exclude infeasible regions)
3. Feasibility-first search (active learning + trust region)
4. Simulation execution (external Motor-CAD/Maxwell)
5. AI result analysis (trends, anomalies, optimization suggestions)
6. Multi-fidelity calibration + confidence grading
7. Experience library enhancement (extract design knowledge)
8. Adaptive next-batch selection (loop back to step 3)
"""
import json
from typing import Dict, List, Optional, Any, Callable
from datetime import datetime
from enum import Enum

from ..services.l0_prescreening import L0PreScreeningEngine
from ..services.feasibility_search import FeasibilityFirstSearch, ParameterRange
from ..services.plan_generator import AIPlanGenerator
from ..services.result_analyst import AIResultAnalyst
from ..services.experience_enhancer import ExperienceEnhancer


class LoopPhase(str, Enum):
    """Phases of the adaptive simulation loop."""
    INIT = "init"
    PLAN_GENERATED = "plan_generated"
    L0_SCREENED = "l0_screened"
    SEARCH_INITIALIZED = "search_initialized"
    BATCH_SELECTED = "batch_selected"
    SIMULATION_RUNNING = "simulation_running"
    RESULTS_ANALYZED = "results_analyzed"
    EXPERIENCE_UPDATED = "experience_updated"
    CONVERGED = "converged"
    BUDGET_EXHAUSTED = "budget_exhausted"
    COMPLETED = "completed"


class AdaptiveLoop:
    """Complete adaptive simulation closed loop.

    This orchestrator ties together all P3-M1 through P3-M5 components
    into a single, stateful simulation optimization loop.
    """

    def __init__(
        self,
        loop_id: Optional[str] = None,
        user_requirement: Optional[str] = None,
        total_budget: int = 80,
        batch_size: int = 4,
    ):
        self.loop_id = loop_id or f"loop_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        self.user_requirement = user_requirement
        self.total_budget = total_budget
        self.batch_size = batch_size
        self.phase = LoopPhase.INIT

        # Components
        self.l0_engine = L0PreScreeningEngine()
        self.plan_generator = AIPlanGenerator(l0_engine=self.l0_engine)
        self.result_analyst = AIResultAnalyst()
        self.experience_enhancer = ExperienceEnhancer()
        self.search: Optional[FeasibilityFirstSearch] = None

        # State
        self.plan: Optional[Dict[str, Any]] = None
        self.plan_validation: Optional[Dict[str, Any]] = None
        self.all_results: List[Dict[str, Any]] = []
        self.latest_analysis: Optional[Dict[str, Any]] = None
        self.latest_insights: Optional[Dict[str, Any]] = None
        self.history: List[Dict[str, Any]] = []
        self.created_at = datetime.now().isoformat()
        self.updated_at = datetime.now().isoformat()

    def generate_plan(self, user_requirement: Optional[str] = None) -> Dict[str, Any]:
        """Step 1: Generate simulation plan from natural language.

        Args:
            user_requirement: Natural language requirement (uses stored if None)

        Returns:
            Generated plan with validation
        """
        if user_requirement:
            self.user_requirement = user_requirement

        if not self.user_requirement:
            raise ValueError("No user requirement provided")

        result = self.plan_generator.generate(self.user_requirement)
        self.plan = result.get("plan", {})
        self.plan_validation = result.get("validation", {})
        self.phase = LoopPhase.PLAN_GENERATED
        self._record_history("plan_generated", {"plan_name": self.plan.get("plan_name", "")})

        return {
            "loop_id": self.loop_id,
            "phase": self.phase.value,
            "plan": self.plan,
            "validation": self.plan_validation,
        }

    def initialize_search(self) -> Dict[str, Any]:
        """Step 2-3: Initialize feasibility-first search from plan.

        Converts plan scan variables to search parameters and
        generates initial LHS batch with L0 pre-screening.
        """
        if not self.plan:
            raise RuntimeError("Plan not generated. Call generate_plan() first.")

        # Convert scan variables to ParameterRange
        parameters = []
        for var in self.plan.get("scan_variables", []):
            if isinstance(var, dict) and "name" in var and "min_value" in var and "max_value" in var:
                parameters.append(ParameterRange(
                    name=var["name"],
                    min_value=float(var["min_value"]),
                    max_value=float(var["max_value"]),
                    step=var.get("step"),
                    unit=var.get("unit", ""),
                    description=var.get("description", ""),
                ))

        if not parameters:
            raise RuntimeError("No valid scan variables in plan")

        # Get search strategy from plan
        search_strategy = self.plan.get("search_strategy", {})
        acceptance = self.plan.get("acceptance_criteria", {})

        self.search = FeasibilityFirstSearch(
            parameters=parameters,
            l0_engine=self.l0_engine,
            total_budget=search_strategy.get("max_solver_calls", self.total_budget),
            batch_size=search_strategy.get("batch_size", self.batch_size),
            initial_samples=search_strategy.get("initial_samples", 16),
            objective_metric=acceptance.get("objective_metric", "tavg_nm"),
            objective_direction=acceptance.get("objective_direction", "maximize"),
        )

        # Generate initial batch
        initial_batch = self.search.generate_initial_batch()
        self.phase = LoopPhase.SEARCH_INITIALIZED
        self._record_history("search_initialized", {
            "n_parameters": len(parameters),
            "initial_batch_size": len(initial_batch),
        })

        return {
            "loop_id": self.loop_id,
            "phase": self.phase.value,
            "search_id": self.search.state.run_id,
            "initial_batch": [
                {"id": p.id, "params": p.params, "status": p.status}
                for p in initial_batch
            ],
            "state": self.search.get_state_summary(),
        }

    def get_next_batch(self) -> Dict[str, Any]:
        """Step 3 (loop): Select next batch of points to simulate.

        Uses active learning with trust region refinement.
        """
        if not self.search:
            raise RuntimeError("Search not initialized. Call initialize_search() first.")

        batch = self.search.select_next_batch()
        if not batch:
            if self.search.state.convergence_status == "converged":
                self.phase = LoopPhase.CONVERGED
            else:
                self.phase = LoopPhase.BUDGET_EXHAUSTED
            return {
                "loop_id": self.loop_id,
                "phase": self.phase.value,
                "batch": [],
                "message": f"Search {self.search.state.convergence_status}",
            }

        self.phase = LoopPhase.BATCH_SELECTED
        self._record_history("batch_selected", {
            "batch_id": self.search.state.current_batch - 1,
            "batch_size": len(batch),
        })

        return {
            "loop_id": self.loop_id,
            "phase": self.phase.value,
            "batch_id": self.search.state.current_batch - 1,
            "points": [
                {"id": p.id, "params": p.params, "status": p.status}
                for p in batch
            ],
            "state": self.search.get_state_summary(),
        }

    def report_results(self, point_results: List[Dict[str, Any]]) -> Dict[str, Any]:
        """Step 4-5: Report simulation results and trigger analysis.

        Args:
            point_results: List of {point_id, metrics, status} dicts

        Returns:
            Analysis results and next-step recommendations
        """
        if not self.search:
            raise RuntimeError("Search not initialized")

        # Report each result to search
        for pr in point_results:
            point_id = pr.get("point_id")
            metrics = pr.get("metrics", {})
            status = pr.get("status", "ok")
            if point_id is not None:
                self.search.report_result(point_id, metrics, status)
                # Add to all results
                result_entry = {"point_id": point_id, **metrics}
                self.all_results.append(result_entry)

        # Run AI analysis on accumulated results
        targets = self.plan.get("acceptance_criteria", {}).get("hard_constraints", {})
        target_dict = {}
        if isinstance(targets, list):
            for constraint in targets:
                # Parse simple constraints like "efficiency_pct >= 92"
                if ">=" in constraint:
                    parts = constraint.split(">=")
                    target_dict[parts[0].strip()] = float(parts[1].strip())
                elif "<=" in constraint:
                    parts = constraint.split("<=")
                    target_dict[parts[0].strip()] = float(parts[1].strip())

        self.latest_analysis = self.result_analyst.analyze(
            results=self.all_results,
            targets=target_dict if target_dict else None,
            fidelity="L3",
            scan_parameters=[p.name for p in self.search.parameters] if self.search else None,
        )

        self.phase = LoopPhase.RESULTS_ANALYZED
        self._record_history("results_analyzed", {
            "n_results": len(self.all_results),
            "confidence": self.latest_analysis.get("confidence", {}).get("grade", "?"),
        })

        return {
            "loop_id": self.loop_id,
            "phase": self.phase.value,
            "analysis": self.latest_analysis,
            "search_state": self.search.get_state_summary(),
            "recommendation": self._get_recommendation(),
        }

    def update_experience(self) -> Dict[str, Any]:
        """Step 6: Extract insights and update experience library.

        Returns:
            Extracted insights and experience entry
        """
        if len(self.all_results) < 5:
            return {"message": "Insufficient results for experience extraction (need >= 5)"}

        self.latest_insights = self.experience_enhancer.extract_insights(
            results=self.all_results,
            project_context={
                "topology": self.plan.get("topology", ""),
                "plan_name": self.plan.get("plan_name", ""),
            },
        )

        # Generate experience entry
        experience_entry = self.experience_enhancer.generate_experience_entry(
            insights=self.latest_insights,
            project_name=self.plan.get("plan_name", "unknown"),
            topology=self.plan.get("topology", "unknown"),
        )

        self.phase = LoopPhase.EXPERIENCE_UPDATED
        self._record_history("experience_updated", {
            "n_rules": len(experience_entry.get("design_rules", [])),
        })

        return {
            "loop_id": self.loop_id,
            "phase": self.phase.value,
            "insights": self.latest_insights,
            "experience_entry": experience_entry,
        }

    def check_completion(self) -> Dict[str, Any]:
        """Check if loop should terminate.

        Returns:
            Completion status and reason
        """
        if not self.search:
            return {"completed": False, "reason": "search_not_initialized"}

        state = self.search.get_state_summary()

        if state["convergence_status"] == "converged":
            self.phase = LoopPhase.CONVERGED
            return {"completed": True, "reason": "converged", "state": state}

        if state["remaining_budget"] <= 0:
            self.phase = LoopPhase.BUDGET_EXHAUSTED
            return {"completed": True, "reason": "budget_exhausted", "state": state}

        return {"completed": False, "reason": "continue", "state": state}

    def _get_recommendation(self) -> str:
        """Generate next-step recommendation based on current state."""
        if not self.search:
            return "Initialize search first"

        state = self.search.get_state_summary()
        analysis = self.latest_analysis or {}
        confidence = analysis.get("confidence", {}).get("grade", "?")

        if state["convergence_status"] == "converged":
            return f"Search converged. Best: {state['best_objective_value']}. Confidence: {confidence}."

        if state["remaining_budget"] <= 0:
            return "Budget exhausted."

        if confidence in ("D",):
            return f"Low confidence ({confidence}). Consider increasing sample count or fidelity level."

        if state["trust_region_active"]:
            return f"Trust region active (radius={state['trust_region_radius']:.3f}). Continue exploitation. Best: {state['best_objective_value']}"

        return f"Exploration phase. {state['feasible_points']} feasible points found. Continue sampling."

    def _record_history(self, event: str, data: Dict[str, Any]):
        """Record a history event."""
        self.history.append({
            "event": event,
            "timestamp": datetime.now().isoformat(),
            "phase": self.phase.value,
            "data": data,
        })
        self.updated_at = datetime.now().isoformat()

    def get_state(self) -> Dict[str, Any]:
        """Get complete loop state."""
        return {
            "loop_id": self.loop_id,
            "phase": self.phase.value,
            "user_requirement": self.user_requirement,
            "total_budget": self.total_budget,
            "batch_size": self.batch_size,
            "plan_name": self.plan.get("plan_name", "") if self.plan else None,
            "n_results": len(self.all_results),
            "search_state": self.search.get_state_summary() if self.search else None,
            "latest_confidence": self.latest_analysis.get("confidence", {}).get("grade") if self.latest_analysis else None,
            "history": self.history,
            "created_at": self.created_at,
            "updated_at": self.updated_at,
        }


# Global loop registry (in-memory, P3-M5 prototype)
_loops: Dict[str, AdaptiveLoop] = {}


def get_loop(loop_id: str) -> Optional[AdaptiveLoop]:
    """Get an existing adaptive loop by ID."""
    return _loops.get(loop_id)


def create_loop(user_requirement: str, total_budget: int = 80, batch_size: int = 4) -> AdaptiveLoop:
    """Create a new adaptive loop."""
    loop = AdaptiveLoop(
        user_requirement=user_requirement,
        total_budget=total_budget,
        batch_size=batch_size,
    )
    _loops[loop.loop_id] = loop
    return loop


def list_loops() -> List[Dict[str, Any]]:
    """List all active loops."""
    return [
        {"loop_id": lid, "phase": loop.phase.value, "n_results": len(loop.all_results)}
        for lid, loop in _loops.items()
    ]