from typing import Any, Optional, Callable

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


class LinearKernel(Kernel):
    def _compute(self, x: Any, y: Optional[Any] = None):

        if y is None:
            y = self.X

        if getattr(x, "ndim", 1) == 1:
            x = x[None, :]
        if getattr(y, "ndim", 1) == 1:
            y = y[None, :]  # pyright: ignore[reportOptionalSubscript]

        return (x @ y.T).squeeze()

    def apply(self, X: Any, v: Any) -> Any:
        return X @ (self.X.T @ v)

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


class PhiKernel(LinearKernel):
    def __init__(
        self, phi: Callable, X: Optional[Any] = None, **kwargs
    ) -> None:
        self.phi = phi
        debug(
            f"PhiKernel: applying projection function {getattr(phi, '__name__', 'phi')}"
        )
        phi_X = self.phi(X) if X is not None else None

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

    def _compute(self, x: Any, y: Optional[Any] = None):
        x_phi = self.phi(x)
        y_phi = self.phi(y) if y is not None else None

        return super()._compute(x_phi, y_phi)

    def __repr__(self):
        return f"PhiKernel(phi={getattr(self.phi, '__name__', 'phi')}, alpha={self._alpha:.2e})"
