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

"""Feasibility-First Search Framework (P3-M2).

Pure Python implementation (no numpy/scipy/sklearn dependency).
Implements:
- LHS (Latin Hypercube Sampling) for initial space coverage
- Active learning batch selection (distance + uncertainty heuristic)
- Constrained feasibility prioritization
- Local trust region refinement
- Search state management with checkpoint support

Per third-party review: default mode is feasibility-first
(constrained Bayesian / active learning), not fixed full-factorial DoE.
"""
import math
import random
import hashlib
import json
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple, Any, Callable
from datetime import datetime

from .l0_prescreening import L0PreScreeningEngine, FeasibilityReport


@dataclass
class ParameterRange:
    """Definition of a scan parameter range."""
    name: str
    min_value: float
    max_value: float
    step: Optional[float] = None  # if set, values are quantized
    unit: str = ""
    description: str = ""

    def normalize(self, value: float) -> float:
        """Normalize value to [0, 1] range."""
        if self.max_value == self.min_value:
            return 0.5
        return (value - self.min_value) / (self.max_value - self.min_value)

    def denormalize(self, norm: float) -> float:
        """Denormalize from [0, 1] to actual range."""
        value = self.min_value + norm * (self.max_value - self.min_value)
        if self.step is not None and self.step > 0:
            value = round(value / self.step) * self.step
            value = max(self.min_value, min(self.max_value, value))
        return value

    def random_value(self, rng: random.Random) -> float:
        """Generate a random value within range."""
        return self.denormalize(rng.random())


@dataclass
class SearchPoint:
    """A single point in the search space."""
    id: int
    params: Dict[str, float]
    status: str = "pending"  # pending / running / ok / failed / infeasible
    metrics: Dict[str, float] = field(default_factory=dict)
    feasibility_report: Optional[Dict[str, Any]] = None
    surrogate_prediction: Optional[float] = None
    surrogate_uncertainty: Optional[float] = None
    batch_id: int = 0
    created_at: str = ""

    def param_hash(self) -> str:
        """Generate a hash of parameter values for caching/dedup."""
        sorted_params = sorted(self.params.items())
        param_str = json.dumps(sorted_params, sort_keys=True)
        return hashlib.md5(param_str.encode()).hexdigest()[:12]


@dataclass
class SearchState:
    """Complete state of an adaptive search run."""
    run_id: str
    parameters: List[ParameterRange]
    points: List[SearchPoint] = field(default_factory=list)
    current_batch: int = 0
    total_budget: int = 80
    used_budget: int = 0
    batch_size: int = 4
    search_method: str = "active_learning"
    trust_region_active: bool = False
    trust_region_center: Optional[Dict[str, float]] = None
    trust_region_radius: float = 0.3  # fraction of normalized range
    best_feasible_point: Optional[SearchPoint] = None
    best_objective_value: float = float('inf')
    objective_metric: str = "tavg_nm"  # metric to optimize
    objective_direction: str = "maximize"  # maximize / minimize
    convergence_status: str = "searching"  # searching / converged / stalled / budget_exhausted
    history: List[Dict[str, Any]] = field(default_factory=list)
    created_at: str = ""
    updated_at: str = ""

    def get_param_names(self) -> List[str]:
        return [p.name for p in self.parameters]

    def get_param_by_name(self, name: str) -> Optional[ParameterRange]:
        for p in self.parameters:
            if p.name == name:
                return p
        return None

    def get_pending_points(self) -> List[SearchPoint]:
        return [p for p in self.points if p.status == "pending"]

    def get_completed_points(self) -> List[SearchPoint]:
        return [p for p in self.points if p.status in ("ok", "failed", "infeasible")]

    def get_feasible_points(self) -> List[SearchPoint]:
        return [p for p in self.points if p.status == "ok" and p.feasibility_report and p.feasibility_report.get("feasible", False)]

    def remaining_budget(self) -> int:
        return self.total_budget - self.used_budget


class FeasibilityFirstSearch:
    """Feasibility-first adaptive search engine.

    Implements the recommended default path from third-party review:
    L0 pre-screening + initial samples + active learning batch selection
    + local trust region refinement.
    """

    def __init__(
        self,
        parameters: List[ParameterRange],
        l0_engine: Optional[L0PreScreeningEngine] = None,
        total_budget: int = 80,
        batch_size: int = 4,
        initial_samples: int = 16,
        objective_metric: str = "tavg_nm",
        objective_direction: str = "maximize",
        seed: int = 42,
    ):
        self.rng = random.Random(seed)
        self.parameters = parameters
        self.l0_engine = l0_engine or L0PreScreeningEngine()
        self.total_budget = total_budget
        self.batch_size = batch_size
        self.initial_samples = initial_samples
        self.objective_metric = objective_metric
        self.objective_direction = objective_direction
        self.state = SearchState(
            run_id=f"search_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
            parameters=parameters,
            total_budget=total_budget,
            batch_size=batch_size,
            objective_metric=objective_metric,
            objective_direction=objective_direction,
            created_at=datetime.now().isoformat(),
        )

    def latin_hypercube_sample(self, n_samples: int) -> List[Dict[str, float]]:
        """Generate Latin Hypercube Samples in normalized space.

        Pure Python implementation (no numpy dependency).
        """
        n_params = len(self.parameters)
        # Create n_samples intervals per dimension
        intervals = []
        for _ in range(n_params):
            # Shuffle interval assignments
            assignments = list(range(n_samples))
            self.rng.shuffle(assignments)
            intervals.append(assignments)

        samples = []
        for i in range(n_samples):
            sample = {}
            for j, param in enumerate(self.parameters):
                # Random point within assigned interval
                interval_idx = intervals[j][i]
                norm = (interval_idx + self.rng.random()) / n_samples
                sample[param.name] = param.denormalize(norm)
            samples.append(sample)
        return samples

    def _min_distance_to_existing(self, params: Dict[str, float], existing: List[Dict[str, float]]) -> float:
        """Calculate minimum normalized Euclidean distance to existing points."""
        if not existing:
            return float('inf')
        min_dist = float('inf')
        for ex in existing:
            dist_sq = 0.0
            for param in self.parameters:
                v1 = param.normalize(params.get(param.name, param.min_value))
                v2 = param.normalize(ex.get(param.name, param.min_value))
                dist_sq += (v1 - v2) ** 2
            dist = math.sqrt(dist_sq)
            min_dist = min(min_dist, dist)
        return min_dist

    def _is_duplicate(self, params: Dict[str, float], existing_points: List[SearchPoint]) -> bool:
        """Check if parameter set is a duplicate of an existing point."""
        param_hash = hashlib.md5(
            json.dumps(sorted(params.items()), sort_keys=True).encode()
        ).hexdigest()[:12]
        return any(p.param_hash() == param_hash for p in existing_points)

    def generate_initial_batch(self) -> List[SearchPoint]:
        """Generate initial LHS samples, filtered by L0 feasibility.

        Prioritizes feasible points but includes some infeasible for
        boundary learning (active learning needs both classes).
        """
        # Generate more samples than needed, then select diverse subset
        n_candidates = max(self.initial_samples * 3, 20)
        candidates = self.latin_hypercube_sample(n_candidates)

        # L0 pre-screen all candidates
        scored = []
        for params in candidates:
            report = self.l0_engine.evaluate(params)
            scored.append((params, report))

        # Separate feasible and infeasible
        feasible = [(p, r) for p, r in scored if r.feasible]
        infeasible = [(p, r) for p, r in scored if not r.feasible]

        # Select diverse subset: mostly feasible, some infeasible for boundary
        n_feasible = min(len(feasible), int(self.initial_samples * 0.75))
        n_infeasible = min(len(infeasible), self.initial_samples - n_feasible)

        selected = []
        selected_params = []

        # Greedy max-distance selection for feasible points
        feasible.sort(key=lambda x: x[1].pass_rate, reverse=True)
        for params, report in feasible:
            if len(selected) >= n_feasible:
                break
            # Check duplicate against already selected params
            is_dup = any(
                hashlib.md5(json.dumps(sorted(p.items()), sort_keys=True).encode()).hexdigest()[:12] ==
                hashlib.md5(json.dumps(sorted(params.items()), sort_keys=True).encode()).hexdigest()[:12]
                for p, _ in selected
            )
            if not is_dup:
                dist = self._min_distance_to_existing(params, selected_params)
                if dist > 0.1 or len(selected) < 3:
                    selected.append((params, report))
                    selected_params.append(params)

        # Add infeasible points for boundary learning
        for params, report in infeasible:
            if len(selected) >= self.initial_samples:
                break
            is_dup = any(
                hashlib.md5(json.dumps(sorted(p.items()), sort_keys=True).encode()).hexdigest()[:12] ==
                hashlib.md5(json.dumps(sorted(params.items()), sort_keys=True).encode()).hexdigest()[:12]
                for p, _ in selected
            )
            if not is_dup:
                selected.append((params, report))
                selected_params.append(params)

        # Create SearchPoint objects
        points = []
        for i, (params, report) in enumerate(selected):
            point = SearchPoint(
                id=len(self.state.points) + i,
                params=params,
                status="infeasible" if not report.feasible else "pending",
                feasibility_report=report.to_dict(),
                batch_id=0,
                created_at=datetime.now().isoformat(),
            )
            if not report.feasible:
                # Mark as completed (L0 rejected) without using solver budget
                point.status = "infeasible"
            points.append(point)

        self.state.points.extend(points)
        self.state.current_batch = 1
        self.state.used_budget += sum(1 for p in points if p.status == "pending")
        self.state.updated_at = datetime.now().isoformat()

        return [p for p in points if p.status == "pending"]

    def select_next_batch(self) -> List[SearchPoint]:
        """Select next batch of points using active learning.

        Strategy:
        1. If feasible points found, activate trust region around best
        2. Select points balancing:
           - Exploration: high distance from existing points
           - Exploitation: near best feasible point (trust region)
           - Uncertainty: points near feasibility boundary
        3. L0 pre-screen all candidates, reject obviously infeasible
        """
        if self.state.remaining_budget() <= 0:
            self.state.convergence_status = "budget_exhausted"
            return []

        # Consume the pending initial LHS batch first: points generated by
        # generate_initial_batch must be simulated before active learning has
        # any observations to learn from. Mark them dispatched so they are not
        # re-selected and so submit_batch targets exactly this batch.
        pending = self.state.get_pending_points()
        if pending:
            batch = pending[: self.batch_size]
            for p in batch:
                p.status = "dispatched"
            self.state.updated_at = datetime.now().isoformat()
            return batch

        feasible_points = self.state.get_feasible_points()
        completed_params = [p.params for p in self.state.get_completed_points()]

        # Activate trust region if we have enough feasible points
        if len(feasible_points) >= 3 and not self.state.trust_region_active:
            # C4 fix: select anchor by optimization direction
            if self.objective_direction == "minimize":
                best = min(feasible_points, key=lambda p: p.metrics.get(self.objective_metric, float("inf")))
            else:
                best = max(feasible_points, key=lambda p: p.metrics.get(self.objective_metric, 0))
            self.state.trust_region_active = True
            self.state.trust_region_center = best.params
            self.state.best_feasible_point = best
            self.state.history.append({
                "event": "trust_region_activated",
                "center": best.params,
                "batch": self.state.current_batch,
            })

        # Generate candidates
        n_candidates = 50
        candidates = []

        if self.state.trust_region_active and self.state.trust_region_center:
            # 70% from trust region, 30% from global exploration
            n_trust = int(n_candidates * 0.7)
            n_global = n_candidates - n_trust

            # Trust region samples (Gaussian-like around center)
            for _ in range(n_trust):
                params = {}
                for param in self.parameters:
                    center_norm = param.normalize(self.state.trust_region_center[param.name])
                    # Sample with decreasing radius as search progresses
                    radius = self.state.trust_region_radius * max(0.3, 1.0 - self.state.used_budget / self.total_budget)
                    sample_norm = center_norm + self.rng.gauss(0, radius * 0.3)
                    sample_norm = max(0.0, min(1.0, sample_norm))
                    params[param.name] = param.denormalize(sample_norm)
                candidates.append(params)

            # Global exploration samples
            candidates.extend(self.latin_hypercube_sample(n_global))
        else:
            # No trust region yet: full LHS exploration
            candidates = self.latin_hypercube_sample(n_candidates)

        # Score and select candidates
        scored_candidates = []
        for params in candidates:
            if self._is_duplicate(params, self.state.points):
                continue

            # L0 pre-screen
            report = self.l0_engine.evaluate(params)
            if not report.feasible:
                continue  # Skip obviously infeasible

            # Distance score (exploration)
            dist = self._min_distance_to_existing(params, completed_params)
            dist_score = min(dist / math.sqrt(len(self.parameters)), 1.0)

            # Trust region proximity score (exploitation)
            trust_score = 0.0
            if self.state.trust_region_center:
                dist_to_center_sq = 0.0
                for param in self.parameters:
                    v1 = param.normalize(params[param.name])
                    v2 = param.normalize(self.state.trust_region_center[param.name])
                    dist_to_center_sq += (v1 - v2) ** 2
                dist_to_center = math.sqrt(dist_to_center_sq)
                trust_score = max(0, 1.0 - dist_to_center / self.state.trust_region_radius)

            # Combined score (balance exploration and exploitation)
            if self.state.trust_region_active:
                score = 0.4 * dist_score + 0.6 * trust_score
            else:
                score = dist_score

            scored_candidates.append((params, report, score))

        # Sort by score and select top batch_size
        scored_candidates.sort(key=lambda x: x[2], reverse=True)
        batch_size = min(self.batch_size, self.state.remaining_budget(), len(scored_candidates))

        selected = []
        for i in range(batch_size):
            params, report, score = scored_candidates[i]
            point = SearchPoint(
                id=len(self.state.points) + i,
                params=params,
                status="dispatched",
                feasibility_report=report.to_dict(),
                batch_id=self.state.current_batch,
                created_at=datetime.now().isoformat(),
            )
            selected.append(point)

        self.state.points.extend(selected)
        self.state.used_budget += len(selected)
        self.state.current_batch += 1
        self.state.updated_at = datetime.now().isoformat()

        return selected

    def report_result(self, point_id: int, metrics: Dict[str, float], status: str = "ok") -> None:
        """Report simulation result for a point.

        Args:
            point_id: ID of the search point
            metrics: Dictionary of metric values
            status: ok / failed
        """
        for point in self.state.points:
            if point.id == point_id:
                point.metrics = metrics
                point.status = status

                # Update best feasible point
                if status == "ok" and self.objective_metric in metrics:
                    value = metrics[self.objective_metric]
                    is_better = (
                        (self.objective_direction == "maximize" and value > self.state.best_objective_value) or
                        (self.objective_direction == "minimize" and value < self.state.best_objective_value)
                    )
                    if is_better or self.state.best_feasible_point is None:
                        self.state.best_objective_value = value
                        self.state.best_feasible_point = point
                        # C5 fix: migrate trust region center to new best point
                        if self.state.trust_region_active:
                            self.state.trust_region_center = point.params

                # Check convergence
                self._check_convergence()
                break

        self.state.updated_at = datetime.now().isoformat()

    def _check_convergence(self) -> None:
        """Check if search has converged or stalled."""
        feasible = self.state.get_feasible_points()
        if len(feasible) < 5:
            return

        # Check if objective has improved in last N points
        recent = feasible[-10:] if len(feasible) >= 10 else feasible
        if len(recent) >= 5:
            values = [p.metrics.get(self.objective_metric, 0) for p in recent]
            if self.objective_direction == "maximize":
                improvement = max(values) - max(values[:-3]) if len(values) > 3 else 0
            else:
                improvement = min(values[:-3]) - min(values) if len(values) > 3 else 0

            if improvement < 0.001 and self.state.trust_region_active:
                # Shrink trust region
                self.state.trust_region_radius *= 0.7
                if self.state.trust_region_radius < 0.05:
                    self.state.convergence_status = "converged"
                    self.state.history.append({
                        "event": "converged",
                        "reason": "trust_region_shrunk_below_threshold",
                        "batch": self.state.current_batch,
                    })

    def get_state_summary(self) -> Dict[str, Any]:
        """Get a summary of current search state."""
        return {
            "run_id": self.state.run_id,
            "search_method": self.state.search_method,
            "convergence_status": self.state.convergence_status,
            "total_budget": self.state.total_budget,
            "used_budget": self.state.used_budget,
            "remaining_budget": self.state.remaining_budget(),
            "current_batch": self.state.current_batch,
            "total_points": len(self.state.points),
            "pending_points": len(self.state.get_pending_points()),
            "completed_points": len(self.state.get_completed_points()),
            "feasible_points": len(self.state.get_feasible_points()),
            "trust_region_active": self.state.trust_region_active,
            "trust_region_radius": round(self.state.trust_region_radius, 4),
            "best_objective_value": self.state.best_objective_value if self.state.best_objective_value != float('inf') else None,
            "best_point_params": self.state.best_feasible_point.params if self.state.best_feasible_point else None,
            "points_history": [
                {
                    "id": _p.id,
                    "batch_id": _p.batch_id,
                    "params": _p.params,
                    "objective": _p.metrics.get(self.state.objective_metric),
                    "feasible": _p.status == "ok",
                    "status": _p.status,
                }
                for _p in self.state.points
                if _p.status in ("ok", "failed", "infeasible")
                and _p.metrics.get(self.state.objective_metric) is not None
            ],
            "infeasible_points": self._count_by_status("infeasible"),
            "failed_points": self._count_by_status("failed"),
            "batch_summary": self._build_batch_summary(),
            "l0_summary": self._build_l0_summary(),
        }

    def _count_by_status(self, status: str) -> int:
        """Count points with the given status."""
        return sum(1 for _p in self.state.points if _p.status == status)

    def _build_batch_summary(self) -> list:
        """Aggregate per-batch point status distribution.

        Returns a list sorted by batch_id, each entry containing
        batch_id, total, pending/ok/infeasible/failed counts and
        best_objective (respecting objective_direction).
        """
        batches = {}
        direction = self.state.objective_direction
        metric = self.state.objective_metric
        for _p in self.state.points:
            b = batches.setdefault(_p.batch_id, {
                "batch_id": _p.batch_id, "total": 0,
                "pending": 0, "dispatched": 0, "ok": 0, "infeasible": 0, "failed": 0,
                "best_objective": None,
            })
            b["total"] += 1
            if _p.status in ("pending", "dispatched", "ok", "infeasible", "failed"):
                b[_p.status] += 1
            if _p.status == "ok":
                obj = _p.metrics.get(metric)
                if obj is not None:
                    if b["best_objective"] is None:
                        b["best_objective"] = obj
                    elif direction == "maximize":
                        b["best_objective"] = max(b["best_objective"], obj)
                    else:
                        b["best_objective"] = min(b["best_objective"], obj)
        return sorted(batches.values(), key=lambda x: x["batch_id"])

    def _build_l0_summary(self) -> dict:
        """Aggregate L0 pre-screening statistics across all points.

        Returns sampled/feasible/infeasible counts, pass_rate and
        top_infeasible_reasons (name/count/category) from failed
        constraint checks in infeasible points.
        """
        points = self.state.points
        total = len(points)
        infeasible_pts = [p for p in points if p.status == "infeasible"]
        feasible_count = total - len(infeasible_pts)
        reason_counts = {}
        for p in infeasible_pts:
            report = p.feasibility_report or {}
            for r in report.get("results", []):
                if not r.get("passed"):
                    name = r.get("name", "unknown")
                    entry = reason_counts.setdefault(name, {
                        "name": name, "count": 0,
                        "category": r.get("category", ""),
                    })
                    entry["count"] += 1
        top_reasons = sorted(
            reason_counts.values(), key=lambda x: -x["count"]
        )[:5]
        return {
            "sampled": total,
            "feasible": feasible_count,
            "infeasible": len(infeasible_pts),
            "pass_rate": round(feasible_count / total, 4) if total else 0.0,
            "top_infeasible_reasons": top_reasons,
        }

    def export_state(self) -> Dict[str, Any]:
        """Export full search state for checkpointing."""
        return {
            "state": {
                "run_id": self.state.run_id,
                "current_batch": self.state.current_batch,
                "total_budget": self.state.total_budget,
                "used_budget": self.state.used_budget,
                "convergence_status": self.state.convergence_status,
                "trust_region_active": self.state.trust_region_active,
                "trust_region_center": self.state.trust_region_center,
                "trust_region_radius": self.state.trust_region_radius,
                # inf (initial sentinel before any result) is not JSON-compliant;
                # export as None, matching get_state_summary().
                "best_objective_value": (
                    self.state.best_objective_value
                    if self.state.best_objective_value not in (float("inf"), float("-inf"))
                    else None
                ),
                "batch_size": self.state.batch_size,
                "objective_metric": self.state.objective_metric,
                "objective_direction": self.state.objective_direction,
                "search_method": self.state.search_method,
            },
            "parameters": [
                {"name": p.name, "min": p.min_value, "max": p.max_value, "step": p.step, "unit": p.unit}
                for p in self.parameters
            ],
            "points": [
                {
                    "id": p.id,
                    "params": p.params,
                    "status": p.status,
                    "metrics": p.metrics,
                    "batch_id": p.batch_id,
                    "feasible": p.feasibility_report.get("feasible") if p.feasibility_report else None,
                }
                for p in self.state.points
            ],
        }

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

        Args:
            payload: dict returned by export_state().
            l0_engine: optional pre-screening engine (reused if omitted).

        Returns:
            A new FeasibilityFirstSearch with the same parameters and
            replayed points (status/metrics/batch preserved).
        """
        params_data = payload.get("parameters", [])
        parameters = [
            ParameterRange(
                name=p["name"],
                min_value=p["min"],
                max_value=p["max"],
                step=p.get("step"),
                unit=p.get("unit", ""),
                description=p.get("description", ""),
            )
            for p in params_data
        ]
        st = payload.get("state", {})
        search = cls(
            parameters=parameters,
            l0_engine=l0_engine,
            total_budget=st.get("total_budget", 80),
            batch_size=st.get("batch_size", 4),
            objective_metric=st.get("objective_metric", "tavg_nm"),
            objective_direction=st.get("objective_direction", "maximize"),
        )
        s = search.state
        if st.get("run_id"):
            s.run_id = st["run_id"]
        if "current_batch" in st:
            s.current_batch = int(st["current_batch"])
        if "used_budget" in st:
            s.used_budget = int(st["used_budget"])
        if st.get("convergence_status"):
            s.convergence_status = st["convergence_status"]
        if "trust_region_active" in st:
            s.trust_region_active = bool(st["trust_region_active"])
        if st.get("trust_region_center") is not None:
            s.trust_region_center = st["trust_region_center"]
        if "trust_region_radius" in st:
            s.trust_region_radius = float(st["trust_region_radius"])
        # export_state serializes the inf sentinel as None (JSON compliance);
        # restore it to the dataclass default rather than assigning None
        # (None would break the `value > best_objective_value` comparison).
        if "best_objective_value" in st and st["best_objective_value"] is not None:
            s.best_objective_value = st["best_objective_value"]
        if st.get("search_method"):
            s.search_method = st["search_method"]
        for pd in payload.get("points", []):
            pt = SearchPoint(
                id=int(pd["id"]),
                params=dict(pd.get("params") or {}),
                status=pd.get("status", "pending"),
                metrics=dict(pd.get("metrics") or {}),
                batch_id=int(pd.get("batch_id", 0)),
            )
            feas = pd.get("feasible")
            if feas is not None:
                pt.feasibility_report = {"feasible": bool(feas)}
            s.points.append(pt)
        if s.best_feasible_point is None:
            for pt in s.points:
                if pt.status == "ok" and (pt.feasibility_report or {}).get("feasible", True):
                    s.best_feasible_point = pt
                    break
        return search