from typing import Any, Optional

from vicentin.kernels import Kernel
from vicentin.utils import Dispatcher
from vicentin.utils.logging import debug


class RBFKernel(Kernel):
    def __init__(
        self, X: Optional[Any] = None, gamma: Optional[float] = None, **kwargs
    ):
        self.gamma = gamma
        self._auto_gamma = gamma is None

        super().__init__(X, **kwargs)

        self.disp_l2_dist = Dispatcher()
        self.disp_exp = Dispatcher()

        try:
            import numpy as np
            from scipy.spatial.distance import cdist

            def np_l2_dist(x, y):
                if y is None:
                    y = x

                return cdist(x, y, metric="sqeuclidean")

            self.disp_l2_dist.register("numpy", np_l2_dist)
            self.disp_exp.register("numpy", np.exp)
        except ImportError:
            pass

        try:
            import torch

            def torch_l2_dist(x, y):
                if y is None:
                    y = x
                return torch.cdist(x, y, p=2.0) ** 2

            self.disp_l2_dist.register("torch", torch_l2_dist)
            self.disp_exp.register("torch", torch.exp)
        except ImportError:
            pass

    def set_X(self, X: Any):
        super().set_X(X)

        if self._auto_gamma:
            self.gamma = 1 / (X.shape[1] * X.var())
            debug(f"RBFKernel: auto-calculated gamma={self.gamma:.4e}")

    def _compute(self, x: Any, y: Optional[Any] = None) -> Any:
        ctx_l2 = self.disp_l2_dist.detect_backend(x)
        ctx_exp = self.disp_exp.detect_backend(x)

        ctx_l2.verbose = False
        ctx_exp.verbose = False

        if getattr(x, "ndim", 1) == 1:
            x = x[None, :]
        if y is not None and getattr(y, "ndim", 1) == 1:
            y = y[None, :]

        dists = self.disp_l2_dist(ctx_l2, x, y)

        return self.disp_exp(ctx_exp, -self.gamma * dists).squeeze()

    def __repr__(self):
        return f"RBFKernel(gamma={self.gamma:.2e}, alpha={self._alpha:.2e})"
