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

"""Adaptive simulation loop orchestrator (P3-M2).

Bridges the web-side feasibility-first search (FeasibilityFirstSearch) to the
local executor through the task system.

Lifecycle:
  start_loop()    build search from plan parameters -> create the initial
                  adaptive_batch task (task_type="adaptive_batch")
  advance_loop()  when the current batch task is completed, read its results
                  (point_id -> params -> metrics), feed them back into the
                  search via report_result(), run AI analysis, then either
                  converge or create the next batch task.

The orchestrator is pull-driven (advance_loop is invoked by a caller / route
/ scheduler), matching the existing poll-based executor model and keeping the
task system free of new completion hooks.

Loop state is persisted as JSON under output/adaptive_loops/ so a restart can
at least recover loop metadata and the current batch task.

All source is ASCII only.
"""
import json
import os
from datetime import datetime
from typing import Any, Dict, List, Optional

from ..services.feasibility_search import FeasibilityFirstSearch, ParameterRange
from ..services.l0_prescreening import L0PreScreeningEngine
from ..services.task_manager import get_task_manager, TaskManager
from ..services.result_analyst import AIResultAnalyst
from ..config import KIMI_API_KEY

# Valid loop phases (mirror the executor/task vocabulary).
LOOP_PHASE_INIT = "initializing"
LOOP_PHASE_RUNNING = "running"
LOOP_PHASE_CONVERGED = "converged"
LOOP_PHASE_BUDGET_EXHAUSTED = "budget_exhausted"
LOOP_PHASE_FAILED = "failed"


def _loop_state_path(state_dir: str, loop_id: str) -> str:
    return os.path.join(state_dir, "%s_loop.json" % loop_id)


class AdaptiveOrchestrator:
    """Coordinates adaptive search <-> task system <-> local executor."""

    def __init__(
        self,
        task_manager: Optional[TaskManager] = None,
        state_dir: Optional[str] = None,
        l0_engine: Optional[L0PreScreeningEngine] = None,
    ):
        self.tm = task_manager or get_task_manager()
        self.l0_engine = l0_engine or L0PreScreeningEngine()
        self.state_dir = state_dir or os.path.join(
            os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))),
            "output", "adaptive_loops",
        )
        os.makedirs(self.state_dir, exist_ok=True)
        self._analyst = AIResultAnalyst()
        # loop_id -> in-memory search + runtime metadata
        self._searches: Dict[str, FeasibilityFirstSearch] = {}
        self._loops: Dict[str, Dict[str, Any]] = {}
        self._load_state()

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------
    def start_loop(
        self,
        loop_id: str,
        parameters: List[Dict[str, Any]],
        plan_id: Optional[int] = None,
        plan_data: Optional[Dict[str, Any]] = None,
        objective_metric: str = "tavg_nm",
        objective_direction: str = "maximize",
        total_budget: int = 80,
        batch_size: int = 4,
        initial_samples: int = 16,
        seed: int = 42,
    ) -> Dict[str, Any]:
        """Start an adaptive loop from explicit search parameters.

        parameters: [{"name", "min_value", "max_value", "step"?, "unit"?}]
        Returns loop status including the initial batch task_id.
        """
        if loop_id in self._loops:
            raise ValueError("Loop %s already exists" % loop_id)

        search = self._build_search(
            parameters, objective_metric, objective_direction,
            total_budget, batch_size, initial_samples, seed,
        )
        self._searches[loop_id] = search
        meta = {
            "loop_id": loop_id,
            "phase": LOOP_PHASE_INIT,
            "plan_id": plan_id,
            "plan_data": plan_data or {},
            "parameters": parameters,
            "objective_metric": objective_metric,
            "objective_direction": objective_direction,
            "total_budget": total_budget,
            "batch_size": batch_size,
            "initial_samples": initial_samples,
            "current_batch": 0,
            "current_task_id": None,
            "n_results": 0,
            "created_at": datetime.now().isoformat(),
            "updated_at": datetime.now().isoformat(),
        }
        self._loops[loop_id] = meta

        # First batch
        batch = search.generate_initial_batch()
        task = self._create_batch_task(loop_id, search, batch, batch_id=0)
        meta["current_batch"] = 0
        meta["current_task_id"] = task.get("task_id")
        meta["phase"] = LOOP_PHASE_RUNNING
        self._save_state(loop_id)
        return self._loop_view(loop_id, include_task=task)

    def advance_loop(self, loop_id: str) -> Dict[str, Any]:
        """Advance one adaptive step: if the current batch task completed,
        feed results back, analyze, then converge or create next batch."""
        meta = self._loops.get(loop_id)
        if meta is None:
            raise ValueError("Loop %s not found" % loop_id)
        search = self._searches.get(loop_id)
        if search is None:
            raise ValueError("Loop %s has no search (process restarted?)" % loop_id)

        if meta["phase"] in (LOOP_PHASE_CONVERGED, LOOP_PHASE_BUDGET_EXHAUSTED, LOOP_PHASE_FAILED):
            return self._loop_view(loop_id)

        task_id = meta.get("current_task_id")
        if not task_id:
            return self._loop_view(loop_id, message="no current task")

        task = self.tm.get_task(task_id)
        if task is None:
            raise ValueError("Loop %s current task %s missing" % (loop_id, task_id))

        if task.get("status") not in ("completed", "failed", "cancelled"):
            # still running - nothing to do yet
            return self._loop_view(loop_id, message="batch still running")

        # ---- batch finished: pull results ----
        results = self.tm.get_task_results(task_id)
        point_results = (results or {}).get("results", [])
        if task.get("status") == "failed":
            meta["phase"] = LOOP_PHASE_FAILED
            meta["updated_at"] = datetime.now().isoformat()
            self._save_state(loop_id)
            return self._loop_view(loop_id, message="batch task failed")

        # Feed results back into the search by point_id
        n_fed = 0
        for r in point_results:
            pid = r.get("point_id")
            if pid is None:
                continue
            metrics = r.get("metrics") or {}
            if not metrics:
                # metrics may be flattened on the result top level
                for k, v in r.items():
                    if k not in ("point_id", "params", "status", "error", "point_index", "solve_time_s"):
                        if isinstance(v, (int, float)):
                            metrics[k] = float(v)
            status = "ok" if r.get("status") == "OK" else "failed"
            search.report_result(int(pid), metrics, status)
            n_fed += 1
        meta["n_results"] += n_fed

        # Optional AI analysis on accumulated results (best-effort; only
        # when the AI backend is configured, else keep quantitative only).
        try:
            analysis = None
            if KIMI_API_KEY:
                analysis = self._analyst.analyze(
                    results=self._collect_results(search),
                    targets=None,
                    fidelity="L3",
                    scan_parameters=[p.name for p in search.parameters],
                )
            meta["latest_analysis"] = analysis
        except Exception as exc:  # noqa: BLE001 - analysis is best-effort
            meta["latest_analysis_error"] = str(exc)

        # ---- decide next step ----
        state = search.get_state_summary()
        if state.get("convergence_status") == "converged":
            meta["phase"] = LOOP_PHASE_CONVERGED
        elif state.get("remaining_budget", 0) <= 0:
            meta["phase"] = LOOP_PHASE_BUDGET_EXHAUSTED
        else:
            nxt = search.select_next_batch()
            if not nxt:
                meta["phase"] = LOOP_PHASE_CONVERGED
            else:
                next_batch_id = int(meta.get("current_batch", 0)) + 1
                task = self._create_batch_task(loop_id, search, nxt, batch_id=next_batch_id)
                meta["current_batch"] = next_batch_id
                meta["current_task_id"] = task.get("task_id")
                meta["phase"] = LOOP_PHASE_RUNNING

        meta["updated_at"] = datetime.now().isoformat()
        self._save_state(loop_id)
        return self._loop_view(loop_id)

    def get_loop_status(self, loop_id: str) -> Dict[str, Any]:
        if loop_id not in self._loops:
            raise ValueError("Loop %s not found" % loop_id)
        return self._loop_view(loop_id)

    def list_loops(self) -> List[Dict[str, Any]]:
        return [
            {
                "loop_id": m["loop_id"],
                "phase": m["phase"],
                "current_batch": m.get("current_batch"),
                "n_results": m.get("n_results"),
                "updated_at": m.get("updated_at"),
            }
            for m in self._loops.values()
        ]

    # ------------------------------------------------------------------
    # Internals
    # ------------------------------------------------------------------
    def _build_search(
        self, parameters, objective_metric, objective_direction,
        total_budget, batch_size, initial_samples, seed,
    ) -> FeasibilityFirstSearch:
        ranges = []
        for p in parameters:
            ranges.append(ParameterRange(
                name=p["name"],
                min_value=float(p["min_value"]),
                max_value=float(p["max_value"]),
                step=float(p["step"]) if p.get("step") else None,
                unit=p.get("unit", ""),
                description=p.get("description", ""),
            ))
        if not ranges:
            raise ValueError("at least one search parameter required")
        return FeasibilityFirstSearch(
            parameters=ranges,
            l0_engine=self.l0_engine,
            total_budget=int(total_budget),
            batch_size=int(batch_size),
            initial_samples=int(initial_samples),
            objective_metric=objective_metric,
            objective_direction=objective_direction,
            seed=int(seed),
        )

    def _create_batch_task(self, loop_id, search, batch, batch_id: int) -> Dict[str, Any]:
        meta = self._loops[loop_id]
        parameters = []
        point_ids = []
        for p in batch:
            pid = getattr(p, "id", None)
            params = getattr(p, "params", None)
            if params is None and isinstance(p, dict):
                params = p.get("params")
                pid = p.get("point_id", p.get("id"))
            point_ids.append(pid)
            item = dict(params or {})
            item["point_id"] = pid
            parameters.append(item)
        task = self.tm.create_task(
            plan_id=meta.get("plan_id"),
            plan_data=meta.get("plan_data") or {},
            parameters=parameters,
            task_name="adaptive-%s-b%s" % (loop_id, batch_id),
            task_type="adaptive_batch",
            loop_id=loop_id,
            batch_id=batch_id,
            point_ids=point_ids,
            dynamic=True,
        )
        return task

    def _collect_results(self, search) -> List[Dict[str, Any]]:
        """Flatten search completed points to [{point_id, **metrics}]."""
        out = []
        for p in search.state.get_completed_points():
            entry = {"point_id": p.id, "status": p.status}
            entry.update(p.metrics)
            out.append(entry)
        return out

    def _loop_view(self, loop_id, include_task=None, message=None) -> Dict[str, Any]:
        meta = self._loops[loop_id]
        search = self._searches.get(loop_id)
        view = {
            "loop_id": loop_id,
            "phase": meta["phase"],
            "plan_id": meta.get("plan_id"),
            "current_batch": meta.get("current_batch"),
            "current_task_id": meta.get("current_task_id"),
            "n_results": meta.get("n_results", 0),
            "updated_at": meta.get("updated_at"),
            "created_at": meta.get("created_at"),
            "search_state": search.get_state_summary() if search else None,
        }
        if include_task:
            view["batch_task"] = include_task
        if message:
            view["message"] = message
        if meta.get("latest_analysis"):
            view["latest_analysis"] = meta["latest_analysis"]
        return view

    # ------------------------------------------------------------------
    # Persistence (best-effort)
    # ------------------------------------------------------------------
    def _save_state(self, loop_id: str) -> None:
        meta = self._loops.get(loop_id)
        if meta is None:
            return
        search = self._searches.get(loop_id)
        payload = dict(meta)
        try:
            if search is not None and hasattr(search, "export_state"):
                payload["search_export"] = search.export_state()
        except Exception:
            pass
        path = _loop_state_path(self.state_dir, loop_id)
        with open(path, "w", encoding="utf-8") as f:
            json.dump(payload, f, ensure_ascii=False, indent=2)

    def _load_state(self) -> None:
        if not os.path.isdir(self.state_dir):
            return
        for fname in os.listdir(self.state_dir):
            if not fname.endswith("_loop.json"):
                continue
            path = os.path.join(self.state_dir, fname)
            try:
                with open(path, "r", encoding="utf-8") as f:
                    meta = json.load(f)
                if isinstance(meta, dict) and meta.get("loop_id"):
                    # metadata recovered; search object is rebuilt on demand
                    self._loops[meta["loop_id"]] = meta
            except Exception:
                continue


# Module-level singleton (mirrors get_task_manager pattern).
_orchestrator: Optional[AdaptiveOrchestrator] = None


def get_orchestrator() -> AdaptiveOrchestrator:
    global _orchestrator
    if _orchestrator is None:
        _orchestrator = AdaptiveOrchestrator()
    return _orchestrator