"""PHYSBO Gaussian-process surrogate as a regressor.

This class intentionally does not implement ask/tell adaptive learning.  It is a
surrogate regressor for comparing PHYSBO's GP/RFM backend with sklearn GPR and
other regressors through the common tkmlr interface.
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import Optional

import numpy as np

from .base import BaseRegressor, as_1d_array, as_2d_array
from .preprocessing import Standardizer
from .registry import register_model


@dataclass
class PhysBORegressor(BaseRegressor):
    score_mode: str = "EI"
    num_rand_basis: int = 200
    interval: int = 0
    standardize: bool = True
    random_seed: Optional[int] = None
    name: str = "physbo"

    supports_std: bool = True
    supports_score: bool = True

    def __post_init__(self):
        self.physbo_ = None
        self.policy_ = None
        self.scaler_ = None
        self.X_train_ = None
        self.y_train_ = None

    def _import_physbo(self):
        if self.physbo_ is None:
            try:
                import physbo
            except ImportError as exc:
                raise ImportError("PHYSBO is not installed. Install it with: pip install physbo") from exc
            self.physbo_ = physbo
        return self.physbo_

    def fit(self, X, y):
        physbo = self._import_physbo()
        X_np = as_2d_array(X)
        y_np = as_1d_array(y)
        if len(X_np) != len(y_np):
            raise ValueError(f"X and y length mismatch: len(X)={len(X_np)}, len(y)={len(y_np)}")

        self.X_train_ = X_np.copy()
        self.y_train_ = y_np.copy()
        idx_train = np.arange(len(y_np), dtype=int)

        if self.standardize:
            self.scaler_ = Standardizer().fit(X_np)
            X_fit = self.scaler_.transform(X_np)
        else:
            self.scaler_ = None
            X_fit = X_np

        self.policy_ = physbo.search.discrete.policy(test_X=X_fit, initial_data=(idx_train, y_np))
        if self.random_seed is not None:
            self.policy_.set_seed(int(self.random_seed))

        # Build/update the surrogate.  simulator=None makes PHYSBO return actions,
        # but the fitted policy also provides posterior mean/covariance.
        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 _transform_X(self, X):
        X_np = as_2d_array(X)
        if self.scaler_ is not None:
            return self.scaler_.transform(X_np)
        return X_np

    def predict(self, X, return_std: bool = False):
        if self.policy_ is None:
            raise RuntimeError("PhysBORegressor is not fitted.")
        X_eval = self._transform_X(X)
        mean = np.asarray(self.policy_.get_post_fmean(X_eval), dtype=float)
        if not return_std:
            return mean
        var = np.asarray(self.policy_.get_post_fcov(X_eval), dtype=float)
        std = np.sqrt(np.maximum(var, 0.0))
        return mean, std

    def acquisition(self, X, mode: Optional[str] = None):
        if self.policy_ is None:
            raise RuntimeError("PhysBORegressor is not fitted.")
        if mode is None:
            mode = self.score_mode
        X_eval = self._transform_X(X)
        return np.asarray(self.policy_.get_score(mode=mode, xs=X_eval), dtype=float)


@register_model("physbo")
def build_physbo(**params):
    return PhysBORegressor(**params)