qdot/planning/state_machine.py

"""
qdot/planning/state_machine.py
================================
Backtracking State Machine — 6-stage autonomous tuning orchestrator.

Stage ordering (blueprint §3.2 + Phase 2.1 addition):
    BOOTSTRAPPING → COARSE_SURVEY → HYPERSURFACE_SEARCH → CHARGE_ID
                 → NAVIGATION → VERIFICATION → COMPLETE

HYPERSURFACE_SEARCH (new): navigate to the charge boundary found by
COARSE_SURVEY before handing off to CHARGE_ID.  Without this stage the
classification step runs at a voltage that may be far from any charge
feature, producing MISC on every attempt.  This matches Schuff et al.
(2025) and Moon et al. (2020) where boundary-walking precedes classification.
"""

from __future__ import annotations

import time
import numpy as np
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple

from qdot.core.types import BacktrackEvent, TuningStage
from qdot.core.state import ExperimentState


# ---------------------------------------------------------------------------
# Stage result
# ---------------------------------------------------------------------------

@dataclass
class StageResult:
    success: bool
    confidence: float
    reason: str
    measurements_taken: int = 0
    data: Dict = field(default_factory=dict)


# ---------------------------------------------------------------------------
# Stage configurations
# ---------------------------------------------------------------------------

@dataclass
class StageConfig:
    stage: TuningStage
    success_threshold: float
    max_retries: int
    max_backtracks: int
    description: str = ""


STAGE_ORDER: List[TuningStage] = [
    TuningStage.BOOTSTRAPPING,
    TuningStage.COARSE_SURVEY,
    TuningStage.HYPERSURFACE_SEARCH,   # NEW
    TuningStage.CHARGE_ID,
    TuningStage.NAVIGATION,
    TuningStage.VERIFICATION,
    TuningStage.COMPLETE,
]

DEFAULT_STAGE_CONFIGS: Dict[TuningStage, StageConfig] = {
    TuningStage.BOOTSTRAPPING: StageConfig(
        stage=TuningStage.BOOTSTRAPPING,
        success_threshold=0.5,
        max_retries=3,
        max_backtracks=0,
        description="Verify device responds to gates and charge sensor is functional",
    ),
    TuningStage.COARSE_SURVEY: StageConfig(
        stage=TuningStage.COARSE_SURVEY,
        success_threshold=0.2,
        max_retries=3,
        max_backtracks=2,
        description="Locate any Coulomb signal in voltage space",
    ),
    TuningStage.HYPERSURFACE_SEARCH: StageConfig(
        stage=TuningStage.HYPERSURFACE_SEARCH,
        success_threshold=0.4,   # boundary_proximity confidence from hypersurface_result()
        max_retries=3,
        max_backtracks=2,
        description=(
            "Navigate to the charge boundary found by COARSE_SURVEY. "
            "Agent walks along conductance gradient until a charge feature "
            "is visible within the scan window."
        ),
    ),
    TuningStage.CHARGE_ID: StageConfig(
        stage=TuningStage.CHARGE_ID,
        success_threshold=0.35,   # was 0.5; ensemble rarely exceeds 0.5 on 3-class boundary cases
        max_retries=2,
        max_backtracks=2,
        description="Classify current charge region via InspectionAgent",
    ),
    TuningStage.NAVIGATION: StageConfig(
        stage=TuningStage.NAVIGATION,
        success_threshold=0.15,   # was 0.7; MAP-based success — belief just needs to concentrate on (1,1)
        max_retries=20,
        max_backtracks=2,
        description="Navigate to target (1,1) charge state via BO",
    ),
    TuningStage.VERIFICATION: StageConfig(
        stage=TuningStage.VERIFICATION,
        success_threshold=0.7,
        max_retries=2,
        max_backtracks=1,
        description="Confirm (1,1) is stable across repeated measurements",
    ),
}


# ---------------------------------------------------------------------------
# State Machine
# ---------------------------------------------------------------------------

class StateMachine:
    """
    Six-stage backtracking state machine.

    Operates on ExperimentState — reads and writes stage,
    consecutive_backtracks, and backtrack_log via the state's helpers.
    """

    def __init__(
        self,
        state: ExperimentState,
        configs: Optional[Dict[TuningStage, StageConfig]] = None,
    ):
        self.state = state
        self.configs = configs if configs is not None else DEFAULT_STAGE_CONFIGS
        self._retries: Dict[TuningStage, int] = {s: 0 for s in TuningStage}
        self._backtracks_at_stage: Dict[TuningStage, int] = {s: 0 for s in TuningStage}

    def process_result(self, result: StageResult) -> Tuple[TuningStage, str, bool]:
        stage = self.state.stage
        config = self.configs.get(stage)
        if config is None:
            return stage, f"No config for stage {stage.name}", False

        # SUCCESS CHECK before HITL — advance immediately when good enough.
        if result.success and result.confidence >= config.success_threshold:
            new_stage, rationale = self._advance(stage, result)
            return new_stage, rationale, False

        hitl, hitl_reason = self._check_hitl(stage, config)
        if hitl:
            return stage, hitl_reason, True

        if self._retries[stage] >= config.max_retries:
            if stage == TuningStage.BOOTSTRAPPING or config.max_backtracks == 0:
                return stage, f"Retries exhausted at {stage.name} with no backtrack available", True
            new_stage, rationale = self._backtrack(stage, result)
            hitl, hitl_reason = self._check_hitl(new_stage, self.configs.get(new_stage, config))
            return new_stage, rationale, hitl

        self._retries[stage] += 1
        rationale = (
            f"Stage {stage.name} attempt {self._retries[stage]}/{config.max_retries} failed "
            f"(confidence={result.confidence:.2f} < threshold={config.success_threshold}). "
            f"Reason: {result.reason}"
        )
        return stage, rationale, False

    def _advance(self, stage: TuningStage, result: StageResult) -> Tuple[TuningStage, str]:
        idx = STAGE_ORDER.index(stage)
        new_stage = STAGE_ORDER[min(idx + 1, len(STAGE_ORDER) - 1)]
        self._retries[stage] = 0
        self.state.advance_stage(new_stage)
        rationale = (
            f"Stage {stage.name} succeeded (confidence={result.confidence:.2f} "
            f">= threshold={self.configs[stage].success_threshold}). "
            f"Advancing to {new_stage.name}."
        )
        return new_stage, rationale

    def _backtrack(self, stage: TuningStage, result: StageResult) -> Tuple[TuningStage, str]:
        idx = STAGE_ORDER.index(stage)
        prev_stage = STAGE_ORDER[max(idx - 1, 0)]
        event = BacktrackEvent(
            run_id=self.state.run_id,
            step=self.state.step,
            timestamp=time.time(),
            from_stage=stage,
            to_stage=prev_stage,
            reason=result.reason,
            consecutive_backtracks_at_level=self.state.consecutive_backtracks + 1,
            hitl_triggered=False,
        )
        self.state.record_backtrack(event)
        self._backtracks_at_stage[stage] += 1
        self._retries[stage] = 0
        self._retries[prev_stage] = 0
        self.state.stage = prev_stage
        rationale = (
            f"Backtracking from {stage.name} to {prev_stage.name} "
            f"after {self.configs[stage].max_retries} retries. "
            f"Reason: {result.reason}. "
            f"Consecutive backtracks: {self.state.consecutive_backtracks}."
        )
        return prev_stage, rationale

    def _check_hitl(self, stage: TuningStage, config: StageConfig) -> Tuple[bool, str]:
        if self.state.consecutive_backtracks >= 2:
            return True, (
                f"Consecutive backtracks >= 2 at stage {stage.name} "
                f"(count={self.state.consecutive_backtracks}). HITL required."
            )
        n_bt = self._backtracks_at_stage.get(stage, 0)
        if n_bt >= config.max_backtracks and config.max_backtracks > 0:
            return True, (
                f"Stage {stage.name} backtrack limit reached "
                f"({n_bt}/{config.max_backtracks}). HITL required."
            )
        stage_count = sum(1 for s in self.state.backtrack_log
                         if s.from_stage == stage or s.to_stage == stage)
        if stage_count > 5:
            return True, f"Loop detected: stage {stage.name} appeared {stage_count} times."
        return False, ""


# ---------------------------------------------------------------------------
# Stage result factory functions
# ---------------------------------------------------------------------------

def bootstrap_result(device_responds: bool, signal_detected: bool) -> StageResult:
    success = device_responds and signal_detected
    reasons = []
    if not device_responds:
        reasons.append("gates do not modulate current")
    if not signal_detected:
        reasons.append("no charge sensor signal")
    return StageResult(
        success=success,
        confidence=1.0 if success else 0.0,
        reason="Device OK" if success else "; ".join(reasons),
        data={"device_responds": device_responds, "signal_detected": signal_detected},
    )


def survey_result(peak_found: bool, peak_quality: float) -> StageResult:
    return StageResult(
        success=peak_found,
        confidence=float(np.clip(peak_quality, 0.0, 1.0)),
        reason="Coulomb peak found" if peak_found else "No clear Coulomb peak",
        data={"peak_quality": peak_quality},
    )


def hypersurface_result(boundary_found: bool, proximity_confidence: float) -> StageResult:
    """
    Create StageResult for HYPERSURFACE_SEARCH stage.

    Args:
        boundary_found:        True if a charge boundary is visible in the
                               current scan window (SNR check passed).
        proximity_confidence:  Continuous estimate of how close the agent is
                               to the charge boundary, ∈ [0, 1].  Derived from
                               the conductance peak quality at the new voltage.
    """
    return StageResult(
        success=boundary_found,
        confidence=float(np.clip(proximity_confidence, 0.0, 1.0)),
        reason=(
            "Charge boundary located in scan window"
            if boundary_found
            else "Charge boundary not yet visible; continuing gradient walk"
        ),
        data={
            "boundary_found": boundary_found,
            "proximity_confidence": proximity_confidence,
        },
    )


def charge_id_result(
    label: str,
    confidence: float,
    physics_override: bool = False,
) -> StageResult:
    effective = min(0.65, confidence) if physics_override else confidence
    success = label in ("single-dot", "double-dot") and effective > 0.35  # was 0.5
    reason = f"Classified as {label}"
    if physics_override:
        reason += " (physics override: confidence capped at 0.65)"
    return StageResult(
        success=success,
        confidence=effective,
        reason=reason,
        data={"label": label, "raw_confidence": confidence, "physics_override": physics_override},
    )

def navigation_result(target_reached: bool, belief_confidence: float) -> StageResult:
    return StageResult(
        success=target_reached,   # was: target_reached and belief_confidence >= 0.7
        confidence=belief_confidence,
        reason="(1,1) state reached" if target_reached else "Target not yet reached",
        data={"target_reached": target_reached, "belief_confidence": belief_confidence},
    )


def verification_result(
    stable: bool, reproducibility: float, charge_noise: float,
) -> StageResult:
    success = stable and reproducibility > 0.8 and charge_noise < 0.1
    confidence = float(reproducibility * (1.0 - charge_noise))
    return StageResult(
        success=success,
        confidence=confidence,
        reason=f"Reproducibility={reproducibility:.2f}, charge_noise={charge_noise:.3f}",
        data={"stable": stable, "reproducibility": reproducibility, "charge_noise": charge_noise},
    )