from typing import Any, Optional

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


class CosineKernel(Kernel):
    def __init__(self, X: Optional[Any] = None, **kwargs):
        super().__init__(X, **kwargs)

        self.disp_cosine = Dispatcher()

        try:
            import numpy as np

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

                x_norm = x / np.linalg.norm(x, axis=1, keepdims=True)
                y_norm = y / np.linalg.norm(y, axis=1, keepdims=True)

                return x_norm @ y_norm.T

            self.disp_cosine.register("numpy", np_cosine)
        except ImportError:
            pass

        try:
            import torch.nn.functional as F

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

                x_norm = F.normalize(x, p=2, dim=1)
                y_norm = F.normalize(y, p=2, dim=1)

                return x_norm @ y_norm.T

            self.disp_cosine.register("torch", torch_cosine)
        except ImportError:
            pass

    def _compute(self, x: Any, y: Optional[Any] = None) -> Any:
        ctx = self.disp_cosine.detect_backend(x)
        ctx.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, :]

        return self.disp_cosine(ctx, x, y).squeeze()

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