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
from ..services.task_manager import get_task_manager


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", {})
        # The unified plan uses "variables" (converted from the AI's
        # "scan_variables"); downstream steps read scan_variables. Normalize
        # once here so both names resolve.
        if not self.plan.get("scan_variables") and self.plan.get("variables"):
            self.plan["scan_variables"] = self.plan["variables"]
        # Topology + default model: loops carry no project context, so fill
        # from the registry; otherwise simulation would start with no model.
        from src.afmcore.topology import is_supported, default_model_for
        topo = (self.plan.get("topology") or "").strip().upper()
        if not is_supported(topo):
            topo = "SSSR"
        self.plan["topology"] = topo
        if not (self.plan.get("model_path") or "").strip():
            self.plan["model_path"] = default_model_for(topo)
        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. Accept both the unified
        # "variables" list (start/stop/step or explicit values) and the raw
        # "scan_variables" name; the AI may use min_value/max_value instead.
        parameters = []
        scan_vars = self.plan.get("scan_variables") or self.plan.get("variables") or []
        for var in scan_vars:
            if not isinstance(var, dict) or "name" not in var:
                continue
            lo = var.get("min_value", var.get("start"))
            hi = var.get("max_value", var.get("stop"))
            if lo is not None and hi is not None:
                parameters.append(ParameterRange(
                    name=var["name"],
                    min_value=float(lo),
                    max_value=float(hi),
                    step=var.get("step"),
                    unit=var.get("unit", ""),
                    description=var.get("description", ""),
                ))
            elif isinstance(var.get("values"), list) and len(var["values"]) >= 2:
                vals = [float(x) for x in var["values"]]
                parameters.append(ParameterRange(
                    name=var["name"],
                    min_value=min(vals),
                    max_value=max(vals),
                    step=None,
                    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, INCLUDING the point's input params:
                # experience extraction needs inputs+outputs together to do
                # parameter-sensitivity analysis (metrics-only entries left the
                # AI with "no input parameters provided").
                point_params = {}
                for sp in self.search.state.points:
                    if sp.id == point_id:
                        point_params = dict(sp.params or {})
                        break
                result_entry = {"point_id": point_id, **point_params, **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"),
        )

        # Persist into the experience library so later AI plan generation can
        # reuse the knowledge. Previously the entry was only returned in the
        # HTTP response and lost on restart - the loop's knowledge never
        # actually closed into the library.
        case_id = self._persist_experience_case(experience_entry)

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

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

    def _persist_experience_case(self, entry: Dict[str, Any]) -> Optional[int]:
        """Store an extracted experience entry as an ExperienceCase row.

        Maps the AI-insight structure onto the case model: the best feasible
        point carries params/metrics, the summary + design rules form the
        conclusion. Returns the new case id, or None on failure (persistence
        must never break the loop).
        """
        try:
            from ..database import SessionLocal
            from ..models.experience_case import ExperienceCase

            best = None
            if self.search:
                best = self.search.state.best_feasible_point
            rules = entry.get("design_rules") or []
            rule_texts = [
                (r.get("rule") if isinstance(r, dict) else str(r)) for r in rules
            ]
            conclusion = (entry.get("summary") or "") + "\n\n" + "\n".join(
                f"- {t}" for t in rule_texts if t
            )
            tags = ",".join(entry.get("tags") or [])
            with SessionLocal() as db:
                case = ExperienceCase(
                    source_plan_id=self.loop_id,
                    topology=entry.get("topology", ""),
                    model_path=(self.plan or {}).get("model_path", ""),
                    params_json=json.dumps((best.params if best else {}) or {}, ensure_ascii=False),
                    metrics_json=json.dumps((best.metrics if best else {}) or {}, ensure_ascii=False),
                    boundary_json=json.dumps(
                        (self.plan or {}).get("boundary_conditions", {}) or {},
                        ensure_ascii=False,
                    ),
                    conclusion=conclusion.strip(),
                    tags=tags,
                    rating=0,
                )
                db.add(case)
                db.commit()
                db.refresh(case)
                return case.id
        except Exception:
            return None

    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 submit_batch_to_executor(self, task_manager=None) -> Dict[str, Any]:
        """Submit the current pending batch to the local executor.

        Bridges the web-side adaptive loop to the local Motor-CAD executor
        through the task system: the batch points (returned by the last
        generate_initial_batch / get_next_batch call and held in the search
        pending list) are wrapped into a single adaptive_batch task. The
        local executor claims it, runs each point, and reports results back
        through /report-results, which feeds the search and continues the loop.

        Returns:
            {"task_id", "n_points", "batch_id"}
        """
        if not self.search:
            raise RuntimeError("Search not initialized. Call initialize_search() first.")

        # Submit exactly the current batch (points marked dispatched by
        # select_next_batch), not every pending point - otherwise the whole
        # initial LHS pool would be submitted at once, defeating batching.
        current = [p for p in self.search.state.points if p.status == "dispatched"]
        if not current:
            return {"task_id": None, "n_points": 0, "batch_id": None,
                    "message": "no dispatched batch to submit (call next-batch first)"}

        tm = task_manager or get_task_manager()
        # Merge the plan's writable fixed params into every point: the executor
        # runs the flat parameter dict as-is and does NOT apply plan_data
        # fixed_params itself. Params with motorcad_var=None are not writable
        # on this model and must be skipped (writing them fails the point).
        fixed: Dict[str, Any] = {}
        for fp in (self.plan or {}).get("fixed_params") or []:
            if isinstance(fp, dict) and fp.get("motorcad_var"):
                fixed[fp["motorcad_var"]] = fp.get("value")
        point_ids = []
        parameters = []
        for p in current:
            point_ids.append(p.id)
            item = dict(fixed)
            item.update(p.params or {})
            item["point_id"] = p.id
            parameters.append(item)

        task = tm.create_task(
            plan_id=None,
            plan_data=self.plan or {},
            parameters=parameters,
            task_name="adaptive-%s-b%s" % (self.loop_id, self.search.state.current_batch),
            task_type="adaptive_batch",
            loop_id=self.loop_id,
            batch_id=self.search.state.current_batch,
            point_ids=point_ids,
            dynamic=True,
        )

        self.phase = LoopPhase.SIMULATION_RUNNING
        self._record_history("batch_submitted", {
            "task_id": task.get("task_id"),
            "batch_id": self.search.state.current_batch,
            "n_points": len(point_ids),
        })

        return {
            "task_id": task.get("task_id"),
            "n_points": len(point_ids),
            "batch_id": self.search.state.current_batch,
            "message": "Batch submitted. Local executor will claim and run it.",
        }


    def export_state(self) -> Dict[str, Any]:
        """Serialize the loop (plan + results + search) for checkpointing.

        The search is exported through FeasibilityFirstSearch.export_state();
        combined with restore_state() this is the checkpoint/resume path.
        """
        return {
            "loop_id": self.loop_id,
            "user_requirement": self.user_requirement,
            "total_budget": self.total_budget,
            "batch_size": self.batch_size,
            "phase": self.phase.value,
            "plan": self.plan,
            "plan_validation": self.plan_validation,
            "all_results": list(self.all_results),
            "latest_analysis": self.latest_analysis,
            "latest_insights": self.latest_insights,
            "history": list(self.history),
            "created_at": self.created_at,
            "updated_at": self.updated_at,
            "search": self.search.export_state() if self.search else None,
        }

    @classmethod
    def restore_state(cls, payload: Dict[str, Any], l0_engine=None) -> "AdaptiveLoop":
        """Rebuild a loop from export_state() output (checkpoint resume).

        Args:
            payload: dict returned by export_state().
            l0_engine: optional pre-screening engine for the rebuilt search.

        Returns:
            A new AdaptiveLoop with plan/phase/results/history/search restored
            and registered in the in-memory loop registry.
        """
        loop = cls(
            loop_id=payload.get("loop_id"),
            user_requirement=payload.get("user_requirement"),
            total_budget=payload.get("total_budget", 80),
            batch_size=payload.get("batch_size", 4),
        )
        try:
            if payload.get("phase"):
                loop.phase = LoopPhase(payload["phase"])
        except ValueError:
            pass  # unknown phase -> keep init
        loop.plan = payload.get("plan")
        loop.plan_validation = payload.get("plan_validation")
        loop.all_results = list(payload.get("all_results", []))
        loop.latest_analysis = payload.get("latest_analysis")
        loop.latest_insights = payload.get("latest_insights")
        loop.history = list(payload.get("history", []))
        if payload.get("created_at"):
            loop.created_at = payload["created_at"]
        if payload.get("updated_at"):
            loop.updated_at = payload["updated_at"]
        if payload.get("search"):
            loop.search = FeasibilityFirstSearch.import_state(
                payload["search"], l0_engine=l0_engine or loop.l0_engine
            )
        _loops[loop.loop_id] = loop
        return loop


    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()
    ]