from typing import Any, Optional

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

disp_ridge = Dispatcher()

try:
    from .ridge_regression_np import ridge_regression as rr_np

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

try:
    from .ridge_regression_torch import ridge_regression as rr_torch

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


class RidgeRegression:
    def __init__(self, lambda_: float = 0, backend: Optional[str] = None):
        self.lambda_ = lambda_
        self.backend = backend

        self.weights = None
        self.kernel = None
        self.is_primal = False

    def fit(self, X: Any, y: Any, W: Optional[Any] = None) -> "RidgeRegression":
        ctx = disp_ridge.detect_backend(X, self.backend)
        info(f"Fitting Ridge Regression model using backend: {ctx.backend}")

        if isinstance(X, Kernel):
            self.kernel = X
            X_data = self.kernel.X
        else:
            self.kernel = LinearKernel(X)
            X_data = X

        n, d = X_data.shape
        lambda_reg = self.lambda_ * n

        self.is_primal = (d < n) and not isinstance(X, Kernel)

        if self.is_primal:
            debug(f"Using Primal space optimization (n={n}, d={d})")
            X_input = X_data
        else:
            debug(f"Using Dual space (Kernel) optimization (n={n}, d={d})")
            X_input = self.kernel

        self.weights = disp_ridge(
            ctx, X_input, y, W, lambda_reg, self.is_primal
        )
        info("Ridge Regression fit complete.")

        return self

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

        debug(f"Predicting data of shape {X.shape}")
        if self.is_primal:
            return X @ self.weights

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

    def __repr__(self):
        return f"RidgeRegression(lambda_={self.lambda_:.2e}, backend={self.backend!r})"
