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

try:
    from .svm_np import SVM as SVM_np, predict as predict_np

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

try:
    from .svm_torch import SVM as SVM_torch, predict as predict_torch

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


class SVMClassifier:
    def __init__(
        self,
        kernel: Optional[Kernel] = None,
        loss: str = "hinge",
        C: Optional[float] = 1.0,
        lambda_: Optional[float] = None,
        strategy: str = "ovo",
        smo_threshold: int = 10_000,
        backend: Optional[str] = None,
        **kwargs,
    ):
        if kernel is None:
            kernel = LinearKernel()
        self.kernel = kernel

        self.loss = loss
        self.lambda_ = lambda_
        self.C = C
        self.strategy = strategy
        self.smo_threshold = smo_threshold

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

    def fit(self, X: Any, y: Any):
        if self.loss not in ("hinge", "squared_hinge"):
            raise ValueError("loss must be 'hinge' or 'squared_hinge'")

        ctx = disp_svm.detect_backend(X, self._backend)
        X, y = disp_svm.cast_values(ctx, X, y)

        info(
            f"Training SVM with {self.loss} loss. Strategy: {self.strategy}. Backend: {ctx.backend}"
        )

        n = int(X.shape[0])
        current_C = self.C
        current_alpha = self.kernel.alpha

        if current_C is None:
            if self.lambda_ is not None and self.lambda_ > 0:
                current_C = 1 / (self.lambda_ * n)
            else:
                current_C = float("inf")

        try:
            if self.loss == "squared_hinge":
                if current_C != float("inf"):
                    self.kernel.alpha = current_alpha + 1 / (2 * current_C)
                current_C = float("inf")

            self.kernel.set_X(X)

            self._model_data = disp_svm(
                ctx,
                self.kernel,
                y,
                self.lambda_,
                current_C,
                self.strategy,
                self.smo_threshold,
                **self._kwargs,
            )
        finally:
            self.kernel.alpha = current_alpha

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

    def predict(self, X: Any) -> Any:
        """Evaluates the decision function on X."""
        if self._model_data is None or self.kernel 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.smo_threshold
        )

    def __repr__(self):
        repr_parts = []

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

        if self.lambda_ is not None:
            repr_parts.append(f"lambda={self.lambda_:.2e}")

        if self.C is not None:
            repr_parts.append(f"C={self.C:.2e}")

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

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