experiment/designer.py

"""ExperimentDesigner: the main entry point for designing experiments."""

from typing import Callable, Dict, List, Optional, Tuple

import torch
from torch import Tensor

from physics_informed_bo.config import OptimizationConfig, OptimizerBackend
from physics_informed_bo.experiment.parameter_space import ParameterSpace
from physics_informed_bo.models.hybrid_model import HybridSurrogate
from physics_informed_bo.priors.prior_manager import PriorManager
from physics_informed_bo.priors.data_prior import DataPrior
from physics_informed_bo.priors.physics_prior import PhysicsPrior
from physics_informed_bo.optimizers.factory import create_optimizer
from physics_informed_bo.optimizers.base_optimizer import BaseOptimizer


class ExperimentDesigner:
    """High-level API for physics-informed Bayesian experiment design.

    This is the main user-facing class. It orchestrates:
    1. Parameter space definition
    2. Physics and data prior management
    3. Surrogate model selection and fitting
    4. Acquisition function optimization
    5. Experiment suggestion

    Example:
        designer = ExperimentDesigner(
            parameter_space=space,
            physics_fn=arrhenius_model,
            initial_data=(X_init, y_init),
        )

        # Get next experiment suggestions
        next_experiments = designer.suggest(n=3)

        # After running experiments, update with results
        designer.update(X_new, y_new)
    """

    def __init__(
        self,
        parameter_space: ParameterSpace,
        physics_fn: Optional[Callable[[Tensor], Tensor]] = None,
        initial_data: Optional[Tuple[Tensor, Tensor]] = None,
        config: Optional[OptimizationConfig] = None,
        physics_constraints: Optional[List[Dict]] = None,
    ):
        """
        Args:
            parameter_space: The experimental parameter space.
            physics_fn: Optional physics model function.
            initial_data: Optional tuple of (X, y) initial observations.
            config: Optimization configuration. Defaults to sensible settings.
            physics_constraints: Optional list of physics constraint dicts.
        """
        self.parameter_space = parameter_space
        self.config = config or OptimizationConfig()

        # Set up physics prior
        physics_prior = None
        if physics_fn is not None:
            physics_prior = PhysicsPrior(physics_fn=physics_fn)
            if physics_constraints:
                for c in physics_constraints:
                    physics_prior.add_constraint(**c)

        # Set up data prior
        data_prior = DataPrior()
        if initial_data is not None:
            X_init, y_init = initial_data
            if y_init.dim() == 1:
                y_init = y_init.unsqueeze(-1)
            data_prior.X = X_init
            data_prior.y = y_init
            data_prior.feature_names = parameter_space.parameter_names

        # Prior manager
        self.prior_manager = PriorManager(
            physics_prior=physics_prior,
            data_prior=data_prior,
        )

        # Build surrogate
        self._surrogate: Optional[HybridSurrogate] = None
        self._optimizer: Optional[BaseOptimizer] = None
        self._iteration = 0

        # Initialize if we have enough data
        self._initialize()

    def _initialize(self) -> None:
        """Initialize surrogate model and optimizer."""
        try:
            mode = self.prior_manager.recommend_surrogate_mode()
        except ValueError:
            # Not enough data or physics model
            return

        self._surrogate = self.prior_manager.build_surrogate(
            mode=mode,
            kernel="matern",
            noise_variance=self.config.noise_variance,
            device=self.config.device,
        )

        # Set up optimizer
        self._optimizer = create_optimizer(self.config)
        self._optimizer.set_surrogate(self._surrogate)
        self._optimizer.set_bounds(self.parameter_space.bounds)

        if self.prior_manager.physics_prior:
            self._optimizer.set_physics_prior(self.prior_manager.physics_prior)

    def suggest(self, n: int = 1) -> Tensor:
        """Suggest the next n experiments to run.

        If not enough data exists for GP-based suggestion, falls back to:
        1. Physics-guided sampling (if physics model available)
        2. Latin Hypercube sampling (space-filling design)

        Args:
            n: Number of experiments to suggest.

        Returns:
            Tensor of shape (n, d) with suggested parameter values.
        """
        self._iteration += 1

        # Not enough data for BO: use initial design
        if self._surrogate is None or self.prior_manager.data_prior.n_observations < 3:
            return self._initial_design(n)

        # Re-fit surrogate with latest data
        data = self.prior_manager.data_prior
        if data.n_observations >= 3:
            self._surrogate.fit(data.X, data.y)
            self._optimizer.set_surrogate(self._surrogate)

        # Suggest via optimizer
        candidates = self._optimizer.suggest(
            n_candidates=n,
            X_observed=data.X,
            y_observed=data.y,
        )

        return candidates

    def _initial_design(self, n: int) -> Tensor:
        """Generate initial design points when insufficient data for BO.

        Uses physics model to prioritize promising regions if available.
        """
        if self.prior_manager.physics_prior is not None:
            # Sample candidates and pick those with best physics predictions
            n_candidates = max(n * 20, 200)
            candidates = self.parameter_space.sample_latin_hypercube(n_candidates)

            # Filter by physics constraints
            candidates = self.prior_manager.physics_prior.get_feasible_subset(candidates)
            if len(candidates) < n:
                candidates = self.parameter_space.sample_latin_hypercube(n_candidates)

            # Rank by physics model prediction
            with torch.no_grad():
                physics_scores = self.prior_manager.physics_prior.evaluate(candidates)

            # Select top-n diverse points (greedy furthest-point selection)
            selected = self._select_diverse_top_k(candidates, physics_scores, n)
            return selected
        else:
            return self.parameter_space.sample_latin_hypercube(n)

    def _select_diverse_top_k(
        self, X: Tensor, scores: Tensor, k: int, top_fraction: float = 0.3
    ) -> Tensor:
        """Select k diverse points from the top-scoring candidates."""
        # Pre-filter to top fraction
        n_top = max(k * 3, int(len(X) * top_fraction))
        top_idx = scores.argsort(descending=True)[:n_top]
        X_top = X[top_idx]

        # Greedy furthest-point selection for diversity
        selected_idx = [0]
        for _ in range(k - 1):
            dists = torch.cdist(X_top, X_top[selected_idx]).min(dim=1).values
            next_idx = dists.argmax().item()
            selected_idx.append(next_idx)

        return X_top[selected_idx]

    def update(self, X_new: Tensor, y_new: Tensor) -> None:
        """Update the designer with new experimental observations.

        Args:
            X_new: New input observations (n, d).
            y_new: New output observations (n, 1) or (n,).
        """
        if y_new.dim() == 1:
            y_new = y_new.unsqueeze(-1)

        self.prior_manager.update_with_observations(X_new, y_new)

        # Re-initialize if we now have enough data
        if self._surrogate is None and self.prior_manager.data_prior.n_observations >= 3:
            self._initialize()

    def get_best(self, maximize: bool = True) -> Dict:
        """Get the best observation so far."""
        X_best, y_best = self.prior_manager.data_prior.get_best(maximize)
        params = self.parameter_space.to_dict(X_best.unsqueeze(0))[0]
        return {"parameters": params, "objective": float(y_best)}

    def predict(self, X: Tensor) -> Tuple[Tensor, Tensor]:
        """Get surrogate model predictions at X."""
        if self._surrogate is None:
            if self.prior_manager.physics_prior:
                pred = self.prior_manager.physics_prior.evaluate(X)
                return pred.unsqueeze(-1), torch.ones_like(pred.unsqueeze(-1)) * 0.1
            raise RuntimeError("No surrogate model fitted yet.")
        return self._surrogate.predict(X)

    def model_quality(self) -> Dict:
        """Assess current surrogate model quality."""
        if self._surrogate is None:
            return {"status": "no_model"}
        return self._surrogate.physics_model_quality()

    def summary(self) -> Dict:
        """Get a summary of the current optimization state."""
        return {
            "iteration": self._iteration,
            "n_observations": self.prior_manager.data_prior.n_observations,
            "prior_summary": self.prior_manager.summary(),
            "model_quality": self.model_quality(),
            "parameter_space": {
                "dimension": self.parameter_space.dimension,
                "parameters": self.parameter_space.parameter_names,
            },
        }