from typing import Any, Optional
from abc import ABC, abstractmethod

from vicentin.utils import Dispatcher
from vicentin.utils.logging import info, debug, error, warn

disp_kernel_backend = Dispatcher()
disp_eye = Dispatcher()

try:
    from .kernel_np import KernelBackendNumpy
    import numpy as np

    disp_kernel_backend.register("numpy", KernelBackendNumpy)
    disp_eye.register("numpy", lambda n, x: np.eye(n, dtype=x.dtype))
except ModuleNotFoundError:
    warn("NumPy backend not found.")
except Exception:
    error("NumPy backend failed to load.")

try:
    from .kernel_torch import KernelBackendTorch
    import torch

    disp_kernel_backend.register("torch", KernelBackendTorch)
    disp_eye.register(
        "torch", lambda n, x: torch.eye(n, device=x.device, dtype=x.dtype)
    )
except ModuleNotFoundError:
    warn("Torch backend not found.")
except Exception:
    error("Torch backend failed to load.")


class Kernel(ABC):
    def __init__(
        self,
        X: Optional[Any] = None,
        alpha: float = 1e-6,
        chunk_size: Optional[int] = None,
        cache_size: Optional[int] = None,
        backend_verbose: bool = False,
        backend: Optional[str] = None,
    ) -> None:
        super().__init__()

        self.is_dense = cache_size is None

        self._alpha = alpha
        self._chunk_size = chunk_size
        self._cache_size = cache_size
        self._disp_backend = backend
        self._backend_verbose = backend_verbose
        self._trace = None

        self.X = None
        self.N = None

        if X is not None:
            self.set_X(X)
        elif self._disp_backend is not None:
            self._set_backend(None)
        else:
            self.backend = None

    def _set_backend(self, arg: Any):
        ctx = disp_kernel_backend.detect_backend(arg, self._disp_backend)
        X = disp_kernel_backend.cast_values(ctx, arg)

        ctx.debug_level = not self._backend_verbose
        ctx.verbose = False

        debug(f"Initializing {type(self).__name__} with backend: {ctx.backend}")
        self.backend = disp_kernel_backend(
            ctx,
            self._compute,
            X,
            self._alpha,
            self._chunk_size,
            self._cache_size,
        )

        self.is_dense = self.backend.cache_size >= self.backend.N

    def set_X(self, X: Any):
        if X is None:
            raise ValueError("X is None.")

        self._set_backend(X)

        self.X = self.backend.X
        self.N = self.backend.N
        self._trace = None

        info(
            f"Kernel bound to data with N={self.N}. Mode: {'Dense' if self.is_dense else 'Chunked'}"
        )

    def center(self):
        return CenteredKernel(self)

    @abstractmethod
    def _compute(self, x: Any, y: Optional[Any] = None) -> Any:
        pass

    def __setattr__(self, name: str, value: Any) -> None:
        if (
            type(value).__name__ == "Parameter"
            and "torch" in type(value).__module__
        ):
            if hasattr(self, "backend") and hasattr(
                self.backend, "register_parameter"
            ):
                self.backend.register_parameter(name, value)

        super().__setattr__(name, value)

    def apply(self, X: Any, v: Any) -> Any:
        if self.backend is None:
            warn("Kernel: Setting backend based on X.")
            self._set_backend(X)

        return self.backend.apply_v(X, v)

    def trace(self) -> float:
        if self.backend is None:
            raise RuntimeError("Backend not initialized")

        if self._trace is None:
            debug(f"Computing trace for {type(self).__name__}...")
            self._trace = self.backend.trace()
            debug(f"Trace computed: {self._trace:.4f}")

        return self._trace

    @property
    def alpha(self) -> float:
        return self._alpha

    @alpha.setter
    def alpha(self, value: float):
        self._alpha = float(value)

        if self.backend is not None:
            self.backend.alpha = self._alpha

            self.backend._full_K = None
            self.backend._cache.clear()
            self._trace = None

    @property
    def T(self):
        return self

    def __getitem__(self, key) -> Any:
        if self.backend is None:
            raise RuntimeError("Backend not initialized.")

        return self.backend[key]

    def __matmul__(self, v) -> Any:
        if self.backend is None:
            warn("Kernel: Setting backend based on v.")
            self._set_backend(v)

        return self.backend @ v

    def __call__(self, x: Any, y: Optional[Any] = None) -> Any:
        if self.backend is None:
            warn("Kernel: Setting backend based on x.")
            self._set_backend(x)

        return self.backend(x, y)  # pyright: ignore[reportOptionalCall]

    def __add__(self, other):
        debug(
            f"Composing SumKernel: {type(self).__name__} + {type(other).__name__}"
        )
        if isinstance(other, Kernel):
            return SumKernel(self, other)
        elif isinstance(other, (int, float)):
            if self.backend is None:
                warn("Kernel: Setting backend based on other.")
                self._set_backend(other)

            return SumKernel(self, ConstantKernel(self.backend.X, other))

        return SumKernel(self, MatrixKernel(other))

    def __radd__(self, other):
        return self.__add__(other)

    def __mul__(self, other):
        if isinstance(other, Kernel):
            debug(
                f"Composing ProductKernel: {type(self).__name__} * {type(other).__name__}"
            )
            return ProductKernel(self, other)

        debug(f"Scaling {type(self).__name__} by {other}")
        return ScaledKernel(self, other)

    def __pow__(self, exponent):
        debug(f"Raising {type(self).__name__} to power {exponent}")
        return PowerKernel(self, exponent)

    def __rmul__(self, other):
        return self.__mul__(other)

    def __truediv__(self, other):
        if isinstance(other, (int, float)):
            return ScaledKernel(self, float(1 / other))

        return NotImplemented

    @property
    def shape(self):
        return (self.N, self.N)

    def __repr__(self):
        return f"{self.__class__.__name__}(alpha={self._alpha:.2e}, chunk_size={self._chunk_size}, backend={self._disp_backend})"


class SumKernel(Kernel):
    def __init__(self, k1: Kernel, k2: Kernel, **kwargs):
        self.k1 = k1
        self.k2 = k2

        kwargs.setdefault("chunk_size", k1._chunk_size)
        kwargs["alpha"] = 0

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

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

        self.k1.set_X(self.X)
        self.k2.set_X(self.X)

    def _compute(self, x: Any, y: Optional[Any] = None):
        return self.k1(x, y) + self.k2(x, y)

    def __repr__(self):
        return f"SumKernel(k1={self.k1!r}, k2={self.k2!r})"


class ProductKernel(Kernel):
    def __init__(self, k1: Kernel, k2: Kernel, **kwargs):
        self.k1 = k1
        self.k2 = k2

        kwargs.setdefault("chunk_size", k1._chunk_size)
        kwargs["alpha"] = 0

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

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

        self.k1.set_X(self.X)
        self.k2.set_X(self.X)

    def _compute(self, x: Any, y: Optional[Any] = None):
        return self.k1(x, y) * self.k2(x, y)

    def __repr__(self):
        return f"ProductKernel(k1={self.k1!r}, k2={self.k2!r})"


class ConstantKernel(Kernel):
    def __init__(self, X: Any, c: float, **kwargs):
        super().__init__(X, alpha=0, **kwargs)
        self.c = float(c)

        self.disp_full = Dispatcher()

        try:
            import numpy as np

            self.disp_full.register("numpy", np.full)
        except ImportError:
            pass

        try:
            import torch

            self.disp_full.register("torch", torch.full)
        except ImportError:
            pass

    def _compute(self, x: Any, y: Optional[Any] = None):
        ctx = self.disp_full.detect_backend(x)
        ctx.verbose = False

        n_x = x.shape[0] if getattr(x, "ndim", 0) > 0 else 1
        n_y = y.shape[0] if y is not None else n_x

        return self.disp_full(ctx, (n_x, n_y), self.c)

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


class ScaledKernel(Kernel):
    def __init__(self, K: Kernel, c: float, **kwargs):
        self.K = K
        self.c = c

        kwargs["alpha"] = 0

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

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

    def _compute(self, x: Any, y: Optional[Any] = None):
        return self.c * self.K(x, y)

    def __repr__(self):
        return f"ScaledKernel(K={self.K!r}, c={self.c})"


class PowerKernel(Kernel):
    def __init__(self, K: Kernel, d: float, **kwargs):
        self.K = K
        self.d = d

        kwargs["alpha"] = 0

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

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

    def _compute(self, x: Any, y: Optional[Any] = None):
        return self.K(x, y) ** self.d

    def __repr__(self):
        return f"PowerKernel(K={self.K!r}, d={self.d})"


class CenteredKernel(Kernel):
    def __init__(self, K: Kernel, **kwargs):
        self.K = K
        self.row_means = None
        self.mean = None

        self._ones = None
        self.disp_ones = Dispatcher()

        try:
            import numpy as np

            self.disp_ones.register(
                "numpy", lambda n, ref: np.ones((n, 1), dtype=ref.dtype)
            )
        except ImportError:
            pass

        try:
            import torch

            self.disp_ones.register(
                "torch",
                lambda n, ref: torch.ones(
                    (n, 1), device=ref.device, dtype=ref.dtype
                ),
            )
        except ImportError:
            pass

        kwargs["chunk_size"] = None
        kwargs["alpha"] = 0

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

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

        self.row_means = None
        self.mean = None

    @property
    def ones(self) -> Any:
        if self._ones is None:
            ctx = self.disp_ones.detect_backend(self.X)
            ctx.verbose = False
            self._ones = self.disp_ones(ctx, self.N, self.X)

        return self._ones

    def _compute_training_statistics(self):
        if self.N is None:
            raise RuntimeError()

        n = self.N
        debug(f"Centering kernel: computing statistics for N={n}")

        sum_rows = self.K @ self.ones

        self.row_means = sum_rows.reshape(-1, 1) / n
        self.mean = float(self.row_means.mean())

    def _compute(self, x: Any, y: Optional[Any] = None) -> Any:
        if self.row_means is None or self.mean is None:
            self._compute_training_statistics()

        if self.N is None:
            raise RuntimeError()

        K_xy = self.K(x, y)

        if x is self.X:
            x_means = self.row_means
        else:
            x_means = (self.K.apply(x, self.ones) / self.N).reshape(-1, 1)

        if y is None or y is self.X:
            y_means = self.row_means.T
        else:
            y_means = (self.K.apply(y, self.ones) / self.N).reshape(1, -1)

        return K_xy - x_means - y_means + self.mean

    def __repr__(self):
        return f"CenteredKernel(K={self.K!r})"


class MatrixKernel(Kernel):
    def __init__(self, M: Any, **kwargs):
        super().__init__(None, **kwargs)
        self.M = M
        self._index_map = None
        self.disp_slice = Dispatcher()

        try:
            import numpy as np

            self.disp_slice.register("numpy", lambda m, i, j: m[np.ix_(i, j)])
        except ImportError:
            pass

        try:
            self.disp_slice.register("torch", lambda m, i, j: m[i][:, j])
        except ImportError:
            pass

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

        if "torch" in str(type(X)):
            self._index_map = {
                tuple(row.tolist()): i for i, row in enumerate(X)
            }
        else:
            self._index_map = {tuple(row): i for i, row in enumerate(X)}

    def _get_indices(self, data: Any):
        indices = []
        for row in data:
            row_tuple = (
                tuple(row.tolist()) if hasattr(row, "tolist") else tuple(row)
            )
            idx = self._index_map.get(row_tuple)
            if idx is None:
                raise ValueError("Point not found in MatrixKernel index.")
            indices.append(idx)
        return indices

    def _compute(self, x: Any, y: Optional[Any] = None) -> Any:
        ctx = self.disp_slice.detect_backend(x)
        i = self._get_indices(x)

        if y is None:
            return self.M[i]

        j = self._get_indices(y)

        return self.disp_slice(ctx, self.M, i, j)

    def __repr__(self):
        return f"MatrixKernel(M_shape={getattr(self.M, 'shape', 'unknown')}, alpha={self._alpha:.2e})"


class SubsetKernel(Kernel):
    def __init__(self, K, mask):
        self.K = K
        X_sub = K.X[mask]

        super().__init__(
            X=X_sub,
            alpha=K.alpha,
            chunk_size=K._chunk_size,
            cache_size=K._cache_size,
            backend=K._disp_backend,
        )

    def _compute(self, x, y=None):
        return self.K._compute(x, y)
