from typing import Any, Optional

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


class PolynomialKernel(Kernel):
    def __init__(
        self,
        X: Optional[Any] = None,
        degree: int = 3,
        gamma: Optional[float] = None,
        c: float = 0,
        **kwargs,
    ):
        self.d = degree
        self.gamma = gamma
        self.c = c
        self._auto_gamma = gamma is None

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

        self.disp_matmul = Dispatcher()
        self.disp_power = Dispatcher()

        try:
            import numpy as np

            self.disp_matmul.register("numpy", lambda x, y: x @ y.T)
            self.disp_power.register("numpy", np.power)
        except ImportError:
            pass

        try:
            import torch

            self.disp_matmul.register(
                "torch", lambda x, y: torch.matmul(x, y.t())
            )
            self.disp_power.register("torch", torch.pow)
        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"PolynomialKernel: auto-calculated gamma={self.gamma:.4e}")

    def _compute(self, x: Any, y: Optional[Any] = None) -> Any:
        if self.gamma is None:
            raise RuntimeError(
                "gamma is None. Ensure X is set or gamma is provided."
            )

        ctx_matmul = self.disp_matmul.detect_backend(x)
        ctx_power = self.disp_power.detect_backend(x)

        ctx_matmul.verbose = False
        ctx_power.verbose = False

        if getattr(x, "ndim", 1) == 1:
            x = x[None, :]
        if y is None:
            y = self.X
        elif getattr(y, "ndim", 1) == 1:
            y = y[None, :]

        dot_product = self.disp_matmul(ctx_matmul, x, y)
        res = (self.gamma * dot_product) + self.c

        return self.disp_power(ctx_power, res, self.d).squeeze()

    def __repr__(self):
        return f"PolynomialKernel(degree={self.d}, gamma={self.gamma:.2e}, c={self.c}, alpha={self._alpha:.2e})"
