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_logistic = Dispatcher()
disp_predict = Dispatcher()

try:
    from .logistic_regression_np import (
        logistic_regression as lr_np,
        predict as predict_np,
    )

    disp_logistic.register("numpy", lr_np)
    disp_predict.register("numpy", predict_np)
except ModuleNotFoundError:
    warn("NumPy backend not found.")
except Exception:
    error("NumPy backend failed to load.")

try:
    from .logistic_regression_torch import (
        logistic_regression as lr_torch,
        predict as predict_torch,
    )

    disp_logistic.register("torch", lr_torch)
    disp_predict.register("torch", predict_torch)
except ModuleNotFoundError:
    warn("Torch backend not found.")
except Exception:
    error("Torch backend failed to load.")


class LogisticRegression:
    def __init__(
        self,
        kernel: Optional[Kernel] = None,
        lambda_: float = 0.0,
        max_iter: int = 20,
        tol: float = 1e-5,
        strategy: str = "ovo",
        backend: Optional[str] = None,
        **kwargs,
    ):
        if kernel is None:
            kernel = LinearKernel()
        self.kernel = kernel

        self.lambda_ = lambda_
        self.max_iter = max_iter
        self.tol = tol
        self.strategy = strategy

        self._backend = backend
        self._kwargs = kwargs
        self._model_data = None
        self.is_primal = False

    def fit(self, X: Any, y: Any) -> "LogisticRegression":
        ctx = disp_logistic.detect_backend(X, self._backend)
        X, y = disp_logistic.cast_values(ctx, X, y)

        info(
            f"Training Logistic Regression. Strategy: {self.strategy}. Backend: {ctx.backend}"
        )

        n, d = X.shape
        self.is_primal = (d < n) and isinstance(self.kernel, LinearKernel)

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

        self._model_data = disp_logistic(
            ctx,
            X_input,
            y,
            self.lambda_,
            self.max_iter,
            self.tol,
            self.is_primal,
            self.strategy,
            **self._kwargs,
        )

        info(
            f"Training complete. Strategy used: {self._model_data['strategy']}"
        )
        return self

    def predict(self, X: Any) -> Any:
        """Evaluates the decision function on X and returns the predicted class."""
        if self._model_data is None:
            raise RuntimeError("Model is not fitted yet.")

        ctx = disp_predict.detect_backend(X, self._backend)
        X = disp_predict.cast_values(ctx, X)

        debug(f"Predicting {X.shape[0]} samples using {ctx.backend}")

        return disp_predict(
            ctx, self._model_data, self.kernel, X, self.is_primal
        )

    def __repr__(self):
        repr_parts = []

        if self.strategy == "ovo":
            repr_parts.append("One-vs-One")
        elif self.strategy == "ovr":
            repr_parts.append("One-vs-Rest")

        repr_parts.append(f"lambda={self.lambda_:.2e}")
        repr_parts.append(f"max_iter={self.max_iter}")
        repr_parts.append(f"tol={self.tol:.2e}")

        if self._backend is not None:
            repr_parts.append(f"backend={self._backend!r}")

        return f"LogisticRegression(kernel={self.kernel!r}, {', '.join(repr_parts)})"
