from __future__ import annotations

from typing import Any
import numpy as np

from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import ConstantKernel, RBF, WhiteKernel

from ..base import BaseSurrogate, as_2d_array, as_1d_array


class SklearnGPRSurrogate(BaseSurrogate):
    """sklearn GaussianProcessRegressor wrapper.

    Hyperparameters are optimized by maximizing log marginal likelihood during fit(),
    not only by cross validation.
    """

    supports_std = True

    def __init__(
        self,
        kernel=None,
        alpha: float = 1.0e-10,
        normalize_y: bool = True,
        n_restarts_optimizer: int = 10,
        random_state: int | None = None,
        length_scale: float = 1.0,
        noise_level: float | None = None,
    ):
        if kernel is None:
            kernel = ConstantKernel(1.0, (1.0e-3, 1.0e3)) * RBF(
                length_scale=length_scale,
                length_scale_bounds=(1.0e-4, 1.0e4),
            )
            if noise_level is not None:
                kernel = kernel + WhiteKernel(
                    noise_level=noise_level,
                    noise_level_bounds=(1.0e-12, 1.0e1),
                )

        self.kernel = kernel
        self.alpha = alpha
        self.normalize_y = normalize_y
        self.n_restarts_optimizer = n_restarts_optimizer
        self.random_state = random_state
        self.length_scale = length_scale
        self.noise_level = noise_level
        self.model = GaussianProcessRegressor(
            kernel=self.kernel,
            alpha=self.alpha,
            normalize_y=self.normalize_y,
            n_restarts_optimizer=self.n_restarts_optimizer,
            random_state=self.random_state,
        )

    def fit(self, X, y):
        X = as_2d_array(X)
        y = as_1d_array(y)
        self.model.fit(X, y)
        return self

    def predict(self, X, return_std: bool = False):
        X = as_2d_array(X)
        if return_std:
            mean, std = self.model.predict(X, return_std=True)
            return mean, std
        return self.model.predict(X)

    def get_params(self, deep: bool = True) -> dict[str, Any]:
        return {
            "kernel": self.kernel,
            "alpha": self.alpha,
            "normalize_y": self.normalize_y,
            "n_restarts_optimizer": self.n_restarts_optimizer,
            "random_state": self.random_state,
            "length_scale": self.length_scale,
            "noise_level": self.noise_level,
        }