import numpy as np
from sklearn.gaussian_process.kernels import RBF, Matern
from sklearn.gaussian_process.kernels import ConstantKernel as C
from sklearn.gaussian_process.kernels import Kernel, Hyperparameter

class LaplacianKernel(Kernel):
    """A Laplacian kernel compatible with scikit-learn's GaussianProcessRegressor."""
    def __init__(self, length_scale=1.0, length_scale_bounds=(1e-5, 1e5)):
        self.length_scale = length_scale
        self.length_scale_bounds = length_scale_bounds

    @property
    def hyperparameter_length_scale(self):
        return Hyperparameter("length_scale", "numeric", self.length_scale_bounds)

    def __call__(self, X, Y=None, eval_gradient=False):
        Y = X if Y is None else Y
        dists = np.sqrt(((X[:, None, :] - Y[None, :, :]) ** 2).sum(-1))
        K = np.exp(-dists / self.length_scale)

        if eval_gradient:
            if Y is X:
                # Derivative w.r.t. length_scale
                K_gradient = (dists / (self.length_scale ** 2)) * K
                return K, K_gradient[:, :, np.newaxis]
            else:
                return K, np.empty((X.shape[0], Y.shape[0], 0))
        else:
            return K

    def diag(self, X):
        return np.ones(X.shape[0])

    def is_stationary(self):
        return True


class KernelFactory:
    """
    Factory for constructing sklearn-compatible Gaussian Process kernels.
    Supports Gaussian (RBF), Laplacian, and Matern kernels.
    """

    @staticmethod
    def create(kernel_name: str):
        """
        Return the kernel object based on the kernel name.

        Args:
            kernel_name (str): One of ['gaussian', 'laplacian', 'matern52']

        Returns:
            A kernel object compatible with sklearn.gaussian_process.GaussianProcessRegressor
        """
        kernel_name = kernel_name.lower()

        if kernel_name == "gaussian":
            return C(1.0, (1e-3, 1e3)) * RBF(length_scale=1.0, length_scale_bounds=(10, 15))

        elif kernel_name == "laplacian":
            return C(1.0, (1e-3, 1e3)) * LaplacianKernel(length_scale=1.0, length_scale_bounds=(50, 1e2))

        elif kernel_name == "matern52":

            return C(1.0, (1e-3, 1e3)) * Matern(length_scale=1.0, length_scale_bounds=(1e-3, 1e3), nu=2.5)

        else:
            raise ValueError(f"Unsupported kernel: {kernel_name}. Choose from ['gaussian', 'laplacian', 'matern52'].")
