from __future__ import annotations

from typing import Any
import numpy as np

from ..base import BaseSurrogate, as_2d_array, as_1d_array


class PhysBOSurrogate(BaseSurrogate):
    """PHYSBO posterior wrapper used as a regression surrogate.

    PHYSBO is naturally designed for discrete candidate search.  This wrapper
    builds a PHYSBO policy using the supplied training points as its test_X.
    In the tkbo CustomOptimizer, fit() is normally called with the observed
    subset and predict() is evaluated for X_candidates.

    Note
    ----
    PHYSBO may be more robust when X passed to predict() is identical to the
    candidate set used at policy construction.  For flexible arbitrary-X GPR,
    use SklearnGPRSurrogate or a future BoTorch surrogate.
    """

    supports_std = True

    def __init__(
        self,
        num_rand_basis: int = 200,
        interval: int = 0,
        score_mode: str = "EI",
        random_seed: int | None = None,
    ):
        self.num_rand_basis = num_rand_basis
        self.interval = interval
        self.score_mode = score_mode
        self.random_seed = random_seed
        self.physbo = None
        self.policy = None

    def _import_physbo(self):
        if self.physbo is None:
            try:
                import physbo
            except ImportError:
                print("IMPORT ERROR: physbo")
                print("Install example: pip install physbo")
                raise
            self.physbo = physbo
        return self.physbo

    def fit(self, X, y):
        physbo = self._import_physbo()
        X = as_2d_array(X)
        y = as_1d_array(y)
        idx = np.arange(len(y), dtype=int)
        self.X_train_ = X.copy()
        self.y_train_ = y.copy()
        self.policy = physbo.search.discrete.policy(test_X=X, initial_data=(idx, y))
        if self.random_seed is not None:
            self.policy.set_seed(int(self.random_seed))

        # Build/update internal posterior. simulator=None returns actions.
        self.policy.bayes_search(
            max_num_probes=0,
            simulator=None,
            score=self.score_mode,
            interval=self.interval,
            num_rand_basis=self.num_rand_basis,
        )
        return self

    def predict(self, X, return_std: bool = False):
        if self.policy is None:
            raise RuntimeError("PhysBOSurrogate is not fitted.")
        X = as_2d_array(X)
        mean = self.policy.get_post_fmean(X)
        if return_std:
            var = self.policy.get_post_fcov(X)
            std = np.sqrt(np.maximum(var, 0.0))
            return mean, std
        return mean

    def acquisition(self, X, mode: str | None = None):
        if self.policy is None:
            raise RuntimeError("PhysBOSurrogate is not fitted.")
        if mode is None:
            mode = self.score_mode
        X = as_2d_array(X)
        return self.policy.get_score(mode=mode, xs=X)

    def get_params(self, deep: bool = True) -> dict[str, Any]:
        return {
            "num_rand_basis": self.num_rand_basis,
            "interval": self.interval,
            "score_mode": self.score_mode,
            "random_seed": self.random_seed,
        }