from typing import Optional, Any

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

disp_pca = Dispatcher()

try:
    from .pca_np import PCA as PCA_np

    disp_pca.register("numpy", PCA_np)
except ModuleNotFoundError:
    warn("NumPy backend not found.")
except Exception:
    error("NumPy backend failed to load.")

try:
    from .pca_torch import PCA as PCA_torch

    disp_pca.register("torch", PCA_torch)
except ModuleNotFoundError:
    warn("Torch backend not found.")
except Exception:
    error("Torch backend failed to load.")


class PCA:
    def __init__(
        self,
        arg: Optional[Any] = None,
        kernel: Optional[Kernel] = None,
        *,
        n_components: Optional[int] = None,
        var: Optional[float] = None,
        backend: Optional[str] = None,
    ):
        if arg is not None:
            if isinstance(arg, float) and 0 <= arg <= 1:
                self.n_components = None
                self.var = arg
            elif isinstance(arg, int):
                self.n_components = arg
                self.var = None
            else:
                raise ValueError(
                    "arg must be a float between 0 and 1 or an int."
                )
        else:
            self.n_components = n_components
            self.var = var

        if self.n_components is not None and self.var is not None:
            raise ValueError(
                "PCA: you must specify one among `n_components` and `var`."
            )

        self._eigenvectors = None
        self._eigenvalues = None
        self._total_variance = None
        self._mean = None

        self.kernel = kernel
        self._is_standard_pca = kernel is None

        self.backend = backend

    @property
    def explained_variance(self) -> Any:
        if self._eigenvalues is None or self._total_variance is None:
            raise RuntimeError("Model was not fit.")

        return self._eigenvalues.sum() / self._total_variance

    @property
    def k(self):
        if self._eigenvalues is None:
            raise RuntimeError("Model was not fit.")

        return len(self._eigenvalues)

    @property
    def eigenvalues(self) -> Any:
        if self._eigenvalues is None:
            raise RuntimeError("Model was not fit.")

        return self._eigenvalues

    @property
    def eigenvectors(self) -> Any:
        if self._eigenvectors is None:
            raise RuntimeError("Model was not fit.")

        return self._eigenvectors

    def fit(self, X: Any) -> Any:
        ctx = disp_pca.detect_backend(X, self.backend)
        info(f"Fitting PCA model using backend: {ctx.backend}")

        if self._is_standard_pca:
            debug("Using Standard PCA (Primal space)")
            N = X.shape[0]

            self._mean = X.mean(0)
            self._total_variance = ((X - self._mean) ** 2).sum() / (N - 1)

            self._eigenvalues, self._eigenvectors = disp_pca(
                ctx, X, self.n_components, self.var, is_standard=True
            )
        else:
            debug(f"Using Kernel PCA with {type(self.kernel).__name__}")

            self.kernel.set_X(X)

            if not isinstance(self.kernel, CenteredKernel):
                self.kernel = self.kernel.center()

            self._total_variance = self.kernel.trace()

            self._eigenvalues, self._eigenvectors = disp_pca(
                ctx, self.kernel, self.n_components, self.var, is_standard=False
            )

        info(
            f"Fit complete. Kept {self.k} components and retained {100 * self.explained_variance:.2f}% of variance."
        )

        return self

    def fit_transform(self, X: Any | Kernel) -> Any:
        self.fit(X)
        return self.transform(X)

    def transform(self, X: Any) -> Any:
        if self._eigenvectors is None:
            raise RuntimeError("Model must be fit before calling transform.")

        debug(f"Transforming data of shape {X.shape}")
        if self._is_standard_pca:
            X_centered = X - self._mean
            return X_centered @ self._eigenvectors

        return self.kernel.apply(X, self._eigenvectors)

    def __repr__(self):
        if self.n_components is not None:
            arg_str = f"n_components={self.n_components}"
        elif self.var is not None:
            arg_str = f"var={self.var:.2e}"
        elif not self._is_standard_pca and self.kernel.N is not None:
            arg_str = f"n_components={self.kernel.N - 1}"
        elif self._is_standard_pca:
            arg_str = "n_components=N"
        else:
            arg_str = "n_components=N-1"

        if isinstance(self.kernel, CenteredKernel):
            k = self.kernel.K
        else:
            k = self.kernel

        if self.backend is None:
            return f"PCA({arg_str}, kernel={k!r})"

        return (
            f"PCA({arg_str}, kernel={self.kernel!r}, backend={self.backend!r})"
        )
