from typing import Callable, Optional
from abc import ABC

import torch
import torch.nn as nn

from vicentin.utils.logging import debug


class KernelBackendTorch(nn.Module, ABC):
    def __init__(
        self,
        func: Callable,
        X: Optional[torch.Tensor] = None,
        alpha: float = 1e-6,
        chunk_size: Optional[int] = None,
        cache_size: int = 1000,
    ):
        super().__init__()

        if X is None:
            X = torch.empty((0, 0))

        self.register_buffer("_X", X)
        self._N = int(X.shape[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) -> torch.Tensor:
        if isinstance(self._X, torch.Tensor):
            return self._X
        raise RuntimeError("Buffer _X is not a tensor.")

    @property
    def N(self) -> int:
        return self._N

    def _get_full_matrix(self) -> torch.Tensor:
        return self.forward(self.X, self.X)

    def forward(
        self, x: torch.Tensor, y: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        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:
                result.diagonal().add_(self.alpha)
        else:
            debug(
                f"Torch Backend: Evaluating {n_x}x{n_y} matrix in chunks of {self.chunk_size}."
            )
            result = torch.zeros((n_x, n_y), device=x.device, 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:
                        K_block.diagonal().add_(self.alpha)

                    result[s_x, s_y] = K_block

        if is_train and self.cache_size >= self.N:
            self._full_K = result

        return result

    def __getitem__(self, key):
        with torch.no_grad():
            if self.cache_size >= self.N:
                return self._get_full_matrix()[key]

            if isinstance(key, int):
                if self.cache_size <= 0:
                    return self.forward(self.X[key : key + 1], self.X).view(-1)

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

                    row = self.forward(self.X[key : key + 1], self.X).view(-1)
                    self._cache[key] = row
                return self._cache[key]

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

            if isinstance(key, tuple):
                row_idx, col_idx = key
                x_eval = self.X[row_idx]
                y_eval = self.X[col_idx]
                if x_eval.ndim == 1:
                    x_eval = x_eval.unsqueeze(0)
                if y_eval.ndim == 1:
                    y_eval = y_eval.unsqueeze(0)
                return self.forward(x_eval, y_eval)

            x_eval = self.X[key]
            if x_eval.ndim == 1:
                x_eval = x_eval.unsqueeze(0)
            return self.forward(x_eval, self.X)

    def __matmul__(self, v: torch.Tensor) -> torch.Tensor:
        if self.chunk_size is None or self.cache_size >= self.N:
            return torch.matmul(self._get_full_matrix(), v)

        debug(
            f"Torch Backend: Matmul ({self.N}x{self.N}) in chunks of {self.chunk_size}."
        )
        D = v.shape[1] if v.ndim > 1 else 1
        v_reshaped = v.view(self.N, D)
        result = torch.zeros(
            (self.N, D), device=self.X.device, 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]
            out_chunk = result[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)

                out_chunk.addmm_(
                    self.func(x_chunk, self.X[s_y]), v_reshaped[s_y]
                )

        out = result + self.alpha * v_reshaped
        return out.view(v.shape)

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

        N_new = X.shape[0]
        D_v = v.shape[1] if v.ndim > 1 else 1
        v_reshaped = v.view(self.N, D_v)

        if self.chunk_size is None:
            result = torch.matmul(self.forward(X), v_reshaped)
            return result.view(N_new, -1) if v.ndim > 1 else result.view(-1)

        debug(
            f"Torch Backend: Applying kernel operator on {X.device} (chunked)."
        )
        result = torch.zeros((N_new, D_v), device=X.device, 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]
            out_chunk = result[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)

                out_chunk.addmm_(
                    self.func(x_chunk, self.X[s_y]), v_reshaped[s_y]
                )

        return result.view(X.shape[0], -1) if v.ndim > 1 else result.view(-1)

    def trace(self):
        if self.chunk_size is None or self.cache_size >= self.N:
            return float(torch.trace(self._get_full_matrix()).item())

        debug("Torch Backend: Computing trace via chunked 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 += (
                torch.trace(self.func(x_chunk, x_chunk)).item()
                + x_chunk.shape[0] * self.alpha
            )

        return float(trace_val)
