qdot/planning/bayesian_opt.py

"""
qdot/planning/bayesian_opt.py
==============================
Multi-fidelity Bayesian Optimisation for voltage navigation.

Proposes ActionProposal objects (from qdot.core.types) using a Gaussian
Process surrogate model with CIM-informed prior mean.

Key types (from qdot.core.types — NOT redefined here):
    ActionProposal — proposed ΔV, passed to SafetyCritic before execution
    BOPoint        — single BO observation (stored in ExperimentState.bo_history)
    VoltagePoint   — (vg1, vg2) coordinate

Acquisition function: Upper Confidence Bound (UCB)
    UCB(V) = μ(V) + β·σ(V)

Blueprint reference: §5.1 (POMDP value computation / BO planner)
"""

from __future__ import annotations

import numpy as np
from scipy.optimize import minimize
from typing import Dict, List, Optional, Tuple

# Phase 0 types — never redefine
from qdot.core.types import ActionProposal, BOPoint, ChargeLabel, VoltagePoint
from qdot.core.state import BeliefState
from qdot.simulator.cim import ConstantInteractionDevice


class GaussianProcess:
    """
    Squared-exponential (RBF) kernel GP with physics-informed prior mean.

    Works in (vg1, vg2) space. All inputs/outputs are native VoltagePoint
    or float scalars — no numpy tuples in the public interface.
    """

    def __init__(
        self,
        length_scale: float = 0.05,
        signal_var: float = 0.5,
        noise_var: float = 0.01,
    ):
        self.length_scale = length_scale
        self.signal_var = signal_var
        self.noise_var = noise_var
        self._X: List[Tuple[float, float]] = []   # (vg1, vg2) pairs
        self._y: List[float] = []
        self._K_inv: Optional[np.ndarray] = None
        self.prior_mean_fn = lambda vg1, vg2: 0.0

    def set_prior_mean(self, fn) -> None:
        self.prior_mean_fn = fn

    def fit(self, bo_history: List[BOPoint]) -> None:
        """Fit GP to BO observation history from ExperimentState.bo_history."""
        if not bo_history:
            return
        self._X = [(p.voltage.vg1, p.voltage.vg2) for p in bo_history]
        self._y = [p.score for p in bo_history]
        n = len(self._X)
        K = np.array([[self._k(self._X[i], self._X[j]) for j in range(n)] for i in range(n)])
        K += self.noise_var * np.eye(n)
        try:
            self._K_inv = np.linalg.inv(K)
        except np.linalg.LinAlgError:
            self._K_inv = np.linalg.inv(K + 1e-6 * np.eye(n))

    def predict(self, vg1: float, vg2: float) -> Tuple[float, float]:
        """Return (mean, variance) at (vg1, vg2)."""
        v_test = (vg1, vg2)
        prior = self.prior_mean_fn(vg1, vg2)
        if not self._X or self._K_inv is None:
            return prior, self.signal_var

        k_vec = np.array([self._k(xi, v_test) for xi in self._X])
        y_adj = np.array([
            self._y[i] - self.prior_mean_fn(self._X[i][0], self._X[i][1])
            for i in range(len(self._y))
        ])
        mean = prior + float(k_vec @ self._K_inv @ y_adj)
        var = max(0.0, float(self._k(v_test, v_test) - k_vec @ self._K_inv @ k_vec))
        return mean, var

    def _k(self, a: Tuple[float, float], b: Tuple[float, float]) -> float:
        dist_sq = (a[0] - b[0]) ** 2 + (a[1] - b[1]) ** 2
        return self.signal_var * np.exp(-dist_sq / (2 * self.length_scale ** 2))


class MultiResBO:
    """
    Multi-fidelity Bayesian Optimisation for voltage navigation.

    Usage:
        bo = MultiResBO(belief=state.belief, voltage_bounds=state.voltage_bounds)
        bo.update(state.bo_history)
        proposal = bo.propose(state.current_voltage, l1_max=state.step_caps["l1_max"])
        # → ActionProposal, pass to SafetyCritic.clip() before applying
    """

    def __init__(
        self,
        belief: BeliefState,
        voltage_bounds: Optional[Dict] = None,
        exploration_weight: float = 2.0,
        device: Optional[ConstantInteractionDevice] = None,
    ):
        """
        Args:
            belief: Current BeliefState (read-only — BO observes but doesn't update it).
            voltage_bounds: From ExperimentState.voltage_bounds.
                            Format: {"vg1": {"min": -1.0, "max": 1.0}, "vg2": ...}
            exploration_weight: UCB β parameter.
            device: CIM device for prior mean function.
        """
        self.belief = belief
        self.voltage_bounds = voltage_bounds or {
            "vg1": {"min": -1.0, "max": 1.0},
            "vg2": {"min": -1.0, "max": 1.0},
        }
        self.exploration_weight = exploration_weight
        self.device = device or ConstantInteractionDevice()

        self.gp = GaussianProcess()
        self._update_gp_prior()

    def update(self, bo_history: List[BOPoint]) -> None:
        """
        Re-fit GP to current BO history and refresh CIM prior.

        Call this after each new BOPoint is added to ExperimentState.bo_history.
        """
        self._update_gp_prior()
        self.gp.fit(bo_history)

    def propose(self, current: VoltagePoint, l1_max: float = 0.10) -> ActionProposal:
        """
        Propose a voltage move using UCB acquisition.

        Args:
            current: Current VoltagePoint from ExperimentState.current_voltage.
            l1_max: Step size cap from ExperimentState.step_caps["l1_max"].

        Returns:
            ActionProposal — pass to SafetyCritic.clip() before applying.
        """
        vg1_bounds = self.voltage_bounds.get("vg1", {"min": -1.0, "max": 1.0})
        vg2_bounds = self.voltage_bounds.get("vg2", {"min": -1.0, "max": 1.0})

        # Search bounds: intersection of global bounds and step cap
        search_bounds = [
            (max(vg1_bounds["min"], current.vg1 - l1_max / 2),
             min(vg1_bounds["max"], current.vg1 + l1_max / 2)),
            (max(vg2_bounds["min"], current.vg2 - l1_max / 2),
             min(vg2_bounds["max"], current.vg2 + l1_max / 2)),
        ]

        # Maximise UCB acquisition via L-BFGS-B
        def neg_ucb(xy: np.ndarray) -> float:
            mu, var = self.gp.predict(xy[0], xy[1])
            return -(mu + self.exploration_weight * np.sqrt(var))

        x0 = np.array([current.vg1, current.vg2])
        result = minimize(neg_ucb, x0, method="L-BFGS-B", bounds=search_bounds)

        if result.success:
            new_vg1, new_vg2 = float(result.x[0]), float(result.x[1])
        else:
            # Fallback: random within bounds
            new_vg1 = float(np.random.uniform(search_bounds[0][0], search_bounds[0][1]))
            new_vg2 = float(np.random.uniform(search_bounds[1][0], search_bounds[1][1]))

        delta = VoltagePoint(vg1=new_vg1 - current.vg1, vg2=new_vg2 - current.vg2)
        expected_new = VoltagePoint(vg1=new_vg1, vg2=new_vg2)

        # Info gain: current uncertainty - expected posterior uncertainty
        _, var_now = self.gp.predict(current.vg1, current.vg2)
        _, var_new = self.gp.predict(new_vg1, new_vg2)
        info_gain = max(0.0, float(var_now - var_new))

        return ActionProposal(
            delta_v=delta,
            expected_new_voltage=expected_new,
            info_gain=info_gain,
        )

    def make_bo_point(
        self,
        voltage: VoltagePoint,
        score: float,
        label: ChargeLabel = ChargeLabel.UNKNOWN,
        confidence: float = 0.0,
        step: int = 0,
    ) -> BOPoint:
        """
        Convenience factory for creating BOPoint objects to add to
        ExperimentState.bo_history. ExperimentState.add_classification()
        already creates BOPoints automatically, but this is available for
        explicit BO-driven observations.
        """
        return BOPoint(
            voltage=voltage,
            score=score,
            label=label,
            confidence=confidence,
            step=step,
        )

    # ------------------------------------------------------------------
    # Private
    # ------------------------------------------------------------------

    def _update_gp_prior(self) -> None:
        """Set GP prior mean to CIM prediction weighted by current belief."""
        belief = self.belief

        def prior_fn(vg1: float, vg2: float) -> float:
            if not belief.charge_probs:
                return self.device.current(vg1, vg2)
            mu = 0.0
            for (n1, n2), prob in belief.charge_probs.items():
                mu += prob * self.device.current_for_state(vg1, vg2, n1, n2)
            return float(mu)

        self.gp.set_prior_mean(prior_fn)