from typing import Callable, Optional
from abc import ABC

import numpy as np

from vicentin.utils.logging import debug


class KernelBackendNumpy(ABC):
    def __init__(
        self,
        func: Callable,
        X: Optional[np.ndarray] = None,
        alpha: float = 1e-6,
        chunk_size: Optional[int] = None,
        cache_size: Optional[int] = None,
    ):
        super().__init__()

        self._X = X if X is not None else np.empty((0, 0))

        self.func = func
        self.alpha = alpha
        self.chunk_size = chunk_size

        self._cache = {}
        self.cache_size = cache_size if cache_size is not None else self.N

        self._full_K = None

    @property
    def X(self) -> np.ndarray:
        return self._X

    @property
    def N(self) -> int:
        return self.X.shape[0]

    def __getitem__(self, key):
        if self.cache_size >= self.N:
            return self(self.X, self.X)[key]

        if isinstance(key, (int, np.integer)):
            if self.cache_size <= 0:
                x = self.X[key : key + 1]
                row = np.ravel(self(x))
                if self.alpha > 0:
                    row[key] += self.alpha
                return row

            if key not in self._cache:
                if len(self._cache) >= self.cache_size:
                    self._cache.pop(next(iter(self._cache)))

                x = self.X[key : key + 1]
                row = np.ravel(self(x))

                if self.alpha > 0:
                    row[key] += self.alpha

                self._cache[key] = row

            return self._cache[key]

        if self.chunk_size is None:
            return self(self.X, self.X)[key]

        if isinstance(key, tuple):
            row_idx, col_idx = key
            x_eval = np.atleast_2d(self.X[row_idx])
            y_eval = np.atleast_2d(self.X[col_idx])
            return self(x_eval, y_eval)

        x_eval = np.atleast_2d(self.X[key])
        return self(x_eval, self.X)

    def __call__(self, x: np.ndarray, y: Optional[np.ndarray] = None):
        is_train = (x is self.X) and (y is None or y is self.X)

        if is_train and self._full_K is not None:
            return self._full_K

        y_eval = y if y is not None else self.X
        n_x, n_y = x.shape[0], y_eval.shape[0]

        add_alpha = self.alpha > 0 and (x is y_eval)

        if self.chunk_size is None:
            result = self.func(x, y_eval)
            if add_alpha:
                np.fill_diagonal(result, result.diagonal() + self.alpha)
        else:
            debug(
                f"NumPy Backend: Evaluating kernel in chunks of {self.chunk_size}."
            )
            result = np.zeros((n_x, n_y), dtype=x.dtype)

            for i in range(0, n_x, self.chunk_size):
                s_x = slice(i, i + self.chunk_size)
                x_chunk = x[s_x]

                for j in range(0, n_y, self.chunk_size):
                    debug(f"Evaluating chunk that starts on ({i}, {j})")

                    s_y = slice(j, j + self.chunk_size)

                    K_block = self.func(x_chunk, y_eval[s_y])

                    if add_alpha and i == j:
                        np.fill_diagonal(
                            K_block, K_block.diagonal() + self.alpha
                        )

                    result[s_x, s_y] = K_block

        if is_train and self.cache_size >= self.N:
            debug(
                f"NumPy Backend: Instantiating full {self.N}x{self.N} kernel matrix."
            )
            self._full_K = result

        return result

    def __matmul__(self, v):
        v = np.asanyarray(v)
        if self.chunk_size is None or self.cache_size >= self.N:
            return self(self.X, self.X) @ v

        debug("NumPy Backend: Executing chunked matrix-vector product.")
        v_reshaped = v.reshape(self.N, -1)
        result = np.zeros((self.N, v_reshaped.shape[1]), dtype=self.X.dtype)

        for i in range(0, self.N, self.chunk_size):
            s_x = slice(i, i + self.chunk_size)
            x_chunk = self.X[s_x]

            for j in range(0, self.N, self.chunk_size):
                debug(f"Evaluating chunk that starts on ({i}, {j})")

                s_y = slice(j, j + self.chunk_size)

                result[s_x] += self.func(x_chunk, self.X[s_y]) @ v_reshaped[s_y]

        return result.reshape(v.shape) + self.alpha * v

    def apply_v(self, X: np.ndarray, v: np.ndarray) -> np.ndarray:
        if X is self.X:
            return self @ v

        v = np.asanyarray(v)
        n_new = X.shape[0]
        D_v = v.shape[1] if v.ndim > 1 else 1
        v_reshaped = v.reshape(self.N, D_v)

        if self.chunk_size is None:
            result = self(X) @ v_reshaped
            return result.reshape(-1, D_v) if v.ndim > 1 else result.ravel()

        debug(
            f"NumPy Backend: Applying kernel operator to {n_new} samples (chunked)."
        )
        result = np.zeros((n_new, D_v), dtype=X.dtype)

        for i in range(0, n_new, self.chunk_size):
            s_x = slice(i, i + self.chunk_size)
            x_chunk = X[s_x]

            for j in range(0, self.N, self.chunk_size):
                debug(f"Evaluating chunk that starts on ({i}, {j})")

                s_y = slice(j, j + self.chunk_size)

                result[s_x] += self.func(x_chunk, self.X[s_y]) @ v_reshaped[s_y]

        return result.reshape(-1, D_v) if v.ndim > 1 else result.ravel()

    def trace(self):
        if self.chunk_size is None or self.cache_size >= self.N:
            return float(np.trace(self(self.X, self.X)))

        debug("NumPy Backend: Computing trace via block diagonals.")
        trace_val = 0.0

        for i in range(0, self.N, self.chunk_size):
            debug(f"Evaluating chunk that starts on {i}")

            s_x = slice(i, i + self.chunk_size)
            x_chunk = self.X[s_x]

            trace_val += (
                np.trace(self.func(x_chunk, x_chunk))
                + x_chunk.shape[0] * self.alpha
            )

        return float(trace_val)

    def _normalize_idx(self, key, limit) -> np.ndarray:
        if isinstance(key, slice):
            return np.arange(*key.indices(limit))
        if isinstance(key, (int, np.integer)):
            return np.array([key if key >= 0 else limit + key])
        return np.asanyarray(key)
