pcb-afm-simulation-system/web/backend/app/routers/search.py

"""Search router - L0 pre-screening and feasibility-first adaptive search (P3-M2)."""
from typing import List, Dict, Any, Optional
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel, Field

from ..services.l0_prescreening import L0PreScreeningEngine
from ..services.feasibility_search import FeasibilityFirstSearch, ParameterRange

router = APIRouter(prefix="/api/search", tags=["Search"])

# Global search instances (in-memory, P3-M2 prototype)
_search_instances: Dict[str, FeasibilityFirstSearch] = {}
_l0_engine = L0PreScreeningEngine()


class L0CheckRequest(BaseModel):
    """Request for L0 feasibility check."""
    params: Dict[str, Any] = Field(..., description="Parameter values to check")


class L0CheckResponse(BaseModel):
    """Response for L0 feasibility check."""
    feasible: bool
    total_checks: int
    passed_checks: int
    failed_checks: int
    pass_rate: float
    results: List[Dict[str, Any]]
    risk_items: List[str]


class SearchCreateRequest(BaseModel):
    """Request to create a new adaptive search."""
    parameters: List[Dict[str, Any]] = Field(..., description="Parameter range definitions")
    total_budget: int = Field(default=80, ge=10, le=500)
    batch_size: int = Field(default=4, ge=1, le=16)
    initial_samples: int = Field(default=16, ge=4, le=100)
    objective_metric: str = Field(default="tavg_nm")
    objective_direction: str = Field(default="maximize")
    seed: int = Field(default=42)


class SearchPointResponse(BaseModel):
    """Response for a search point."""
    id: int
    params: Dict[str, float]
    status: str
    metrics: Dict[str, float] = Field(default_factory=dict)
    batch_id: int
    feasible: Optional[bool] = None


class SearchStateResponse(BaseModel):
    """Response for search state summary."""
    run_id: str
    search_method: str
    convergence_status: str
    total_budget: int
    used_budget: int
    remaining_budget: int
    current_batch: int
    total_points: int
    pending_points: int
    completed_points: int
    feasible_points: int
    trust_region_active: bool
    trust_region_radius: float
    best_objective_value: Optional[float] = None
    best_point_params: Optional[Dict[str, float]] = None


class SearchResultRequest(BaseModel):
    """Request to report simulation result."""
    point_id: int
    metrics: Dict[str, float]
    status: str = Field(default="ok", description="ok / failed")


@router.post("/l0/check", response_model=L0CheckResponse)
def l0_check(request: L0CheckRequest):
    """Run L0 analytic pre-screening on a parameter set."""
    report = _l0_engine.evaluate(request.params)
    data = report.to_dict()
    return L0CheckResponse(**data)


@router.post("/l0/filter")
def l0_filter(param_sets: List[Dict[str, Any]]):
    """Filter a list of parameter sets into feasible and infeasible."""
    feasible, infeasible = _l0_engine.filter_feasible(param_sets)
    return {
        "total": len(param_sets),
        "feasible_count": len(feasible),
        "infeasible_count": len(infeasible),
        "feasible": feasible,
        "infeasible": infeasible,
    }


@router.post("/create")
def create_search(request: SearchCreateRequest):
    """Create a new feasibility-first adaptive search run."""
    try:
        parameters = [
            ParameterRange(
                name=p["name"],
                min_value=float(p["min_value"]),
                max_value=float(p["max_value"]),
                step=p.get("step"),
                unit=p.get("unit", ""),
                description=p.get("description", ""),
            )
            for p in request.parameters
        ]
    except (KeyError, ValueError) as e:
        raise HTTPException(status_code=400, detail=f"Invalid parameter definition: {e}")

    search = FeasibilityFirstSearch(
        parameters=parameters,
        l0_engine=_l0_engine,
        total_budget=request.total_budget,
        batch_size=request.batch_size,
        initial_samples=request.initial_samples,
        objective_metric=request.objective_metric,
        objective_direction=request.objective_direction,
        seed=request.seed,
    )

    # Generate initial batch
    initial_points = search.generate_initial_batch()

    _search_instances[search.state.run_id] = search

    return {
        "run_id": search.state.run_id,
        "initial_points": [
            {"id": p.id, "params": p.params, "status": p.status, "batch_id": p.batch_id}
            for p in initial_points
        ],
        "state": search.get_state_summary(),
    }


@router.get("/{run_id}/state", response_model=SearchStateResponse)
def get_search_state(run_id: str):
    """Get current state of a search run."""
    if run_id not in _search_instances:
        raise HTTPException(status_code=404, detail=f"Search run {run_id} not found")
    search = _search_instances[run_id]
    return SearchStateResponse(**search.get_state_summary())


@router.post("/{run_id}/next-batch")
def next_batch(run_id: str):
    """Select and return the next batch of points to simulate."""
    if run_id not in _search_instances:
        raise HTTPException(status_code=404, detail=f"Search run {run_id} not found")
    search = _search_instances[run_id]

    points = search.select_next_batch()
    if not points:
        return {
            "run_id": run_id,
            "batch": [],
            "convergence_status": search.state.convergence_status,
            "message": "No more points to select (budget exhausted or converged)",
        }

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


@router.post("/{run_id}/report")
def report_result(run_id: str, request: SearchResultRequest):
    """Report simulation result for a point."""
    if run_id not in _search_instances:
        raise HTTPException(status_code=404, detail=f"Search run {run_id} not found")
    search = _search_instances[run_id]
    search.report_result(request.point_id, request.metrics, request.status)
    return {
        "run_id": run_id,
        "point_id": request.point_id,
        "status": "reported",
        "state": search.get_state_summary(),
    }


@router.get("/{run_id}/points")
def get_points(run_id: str, status: Optional[str] = None):
    """Get all points in a search run, optionally filtered by status."""
    if run_id not in _search_instances:
        raise HTTPException(status_code=404, detail=f"Search run {run_id} not found")
    search = _search_instances[run_id]

    points = search.state.points
    if status:
        points = [p for p in points if p.status == status]

    return {
        "run_id": run_id,
        "total": len(points),
        "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 points
        ],
    }


@router.get("/{run_id}/export")
def export_search(run_id: str):
    """Export full search state for checkpointing."""
    if run_id not in _search_instances:
        raise HTTPException(status_code=404, detail=f"Search run {run_id} not found")
    search = _search_instances[run_id]
    return search.export_state()


@router.get("/runs")
def list_runs():
    """List all active search runs."""
    return {
        "runs": [
            {"run_id": run_id, "state": search.get_state_summary()}
            for run_id, search in _search_instances.items()
        ],
    }