from typing import Optional
from itertools import combinations

import torch

from vicentin.kernels import Kernel, LinearKernel, SubsetKernel
from vicentin.optimization.problems import QP
from vicentin.machine_learning.classification.svm import SMO
from vicentin.utils.logging import info, debug


def _get_alpha0(y, C, dtype):
    pos_mask = y == 1
    neg_mask = y == -1
    n_pos = pos_mask.sum()
    n_neg = neg_mask.sum()

    epsilon = 0.01 if C == float("inf") else 0.05 * C

    alpha0 = torch.zeros_like(y, dtype=dtype)
    alpha0[pos_mask] = epsilon

    val_neg = (n_pos * epsilon) / n_neg
    alpha0[neg_mask] = -val_neg

    if C != float("inf") and val_neg >= C:
        scale = (0.9 * C) / val_neg
        alpha0 *= scale

    return alpha0


def _train_binary(X, y, kernel, lambda_, C, smo_threshold, **kwargs):
    n = X.shape[0]
    y = y.flatten().to(X.dtype)

    if C is None:
        C = 1 / (lambda_ * n)

    if n > smo_threshold:
        debug(f"Torch: Training with SMO (n={n})")
        return SMO(X, y, kernel, C, **kwargs)

    debug(f"Torch: Training with QP (n={n}, device={X.device})")
    P = kernel[:, :] if isinstance(kernel, torch.Tensor) else kernel(X, X)
    q = -y

    diag_y = torch.diag(y)
    if C == float("inf"):
        G = -diag_y
        h = torch.zeros(n, device=X.device, dtype=X.dtype)
    else:
        G = torch.vstack((diag_y, -diag_y))
        h = torch.ones(2 * n, device=X.device, dtype=X.dtype)
        h[:n] *= C
        h[n:] *= 0

    A = torch.ones((1, n), device=X.device, dtype=X.dtype)
    b_eq = torch.zeros(1, device=X.device, dtype=X.dtype)

    alpha0 = _get_alpha0(y, C, X.dtype)

    alpha_y, duals = QP(P, q, alpha0, G, h, A, b_eq, return_dual=True, **kwargs)
    b_aux = duals[1]

    alpha = alpha_y * y  # Equivalent to dividing because y is -1 or 1.
    sv_mask = (alpha > 1e-4) & (alpha < C - 1e-4)

    debug(f"L2-SVM: SVs found: {torch.sum(sv_mask)}")

    if torch.any(sv_mask):
        b = float(torch.mean(y[sv_mask] - (P @ alpha_y)[sv_mask]))
    else:
        b = float(b_aux[0])

    return alpha_y, b


def SVM(
    X: torch.Tensor | Kernel,
    y: torch.Tensor,
    lambda_: float = 1,
    C: Optional[float] = None,
    strategy: str = "ovo",
    smo_threshold: int = 10_000,
    **kwargs,
):
    if isinstance(X, Kernel):
        kernel = X
        X = torch.as_tensor(kernel.X)
    else:
        kernel = LinearKernel(X)

    y = y.flatten()
    classes = torch.unique(y).int()
    classes_list = classes.tolist()
    info(f"Torch Classes: {classes_list}")

    if len(classes) == 2:
        y = torch.where(
            y == classes[0],
            torch.tensor(-1, dtype=y.dtype, device=y.device),
            torch.tensor(1, dtype=y.dtype, device=y.device),
        )

        alpha, b = _train_binary(
            X, y, kernel, lambda_, C, smo_threshold, **kwargs
        )

        mask = torch.ones_like(y, dtype=torch.bool)
        return {
            "classes": classes,
            "models": [(alpha, b, mask)],
            "strategy": "binary",
        }

    models = []
    if strategy == "ovo":
        info(
            f"Executing One-vs-One strategy ({len(list(combinations(classes, 2)))} pairs)"
        )

        for c1, c2 in combinations(classes_list, 2):
            debug(f"Pair: {c1} vs {c2}")
            mask = (y == c1) | (y == c2)

            y_bin = torch.where(
                y[mask] == c1,
                torch.tensor(-1, dtype=y.dtype, device=y.device),
                torch.tensor(1, dtype=y.dtype, device=y.device),
            )

            sub_kernel = SubsetKernel(kernel, mask)

            alpha, b = _train_binary(
                sub_kernel.X,
                y_bin,
                sub_kernel,
                lambda_,
                C,
                smo_threshold,
                **kwargs,
            )

            models.append((alpha, b, mask))

    elif strategy == "ovr":
        info(f"Torch: One-vs-Rest ({len(classes_list)} classes)")

        for c in classes_list:
            debug(f"Class: {c} vs Rest")
            mask = torch.ones_like(y, dtype=torch.bool)

            y_bin = torch.where(
                y == c,
                torch.tensor(1, dtype=y.dtype, device=y.device),
                torch.tensor(-1, dtype=y.dtype, device=y.device),
            )

            alpha, b = _train_binary(
                X,
                y_bin,
                kernel,
                lambda_,
                C,
                smo_threshold,
                **kwargs,
            )

            models.append((alpha, b, mask))

    return {"classes": classes, "models": models, "strategy": strategy}


def predict(
    model_data,
    X_train: torch.Tensor | Kernel,
    X_test: torch.Tensor,
    cache: int = 10_000,
):
    classes = model_data["classes"]
    models = model_data["models"]
    strategy = model_data["strategy"]

    if isinstance(X_train, Kernel):
        kernel = X_train
        X_train = torch.as_tensor(kernel.X)
    else:
        kernel = LinearKernel(X_train)

    device = X_test.device

    if strategy == "binary":
        info("Prediction using binary strategy")
        alpha, b, _ = models[0]

        pred = kernel.apply(X_test, alpha) + b

        return torch.where(pred < 0, classes[0], classes[1])

    if strategy == "ovr":
        info("Prediction using OvR strategy")
        scores = torch.empty(
            (X_test.shape[0], len(classes)), device=device, dtype=X_test.dtype
        )

        for i, (alpha, b, _) in enumerate(models):
            debug(f"Running model {i}")
            scores[:, i] = kernel.apply(X_test, alpha) + b

        return classes[torch.argmax(scores, dim=1)]

    if strategy == "ovo":
        info("Prediction using OvO strategy")
        votes = torch.zeros(
            (X_test.shape[0], len(classes)), device=device, dtype=X_test.dtype
        )
        idx = 0

        if max(X_test.shape[0], kernel.N) > cache:
            K = None
        else:
            K = kernel(X_test)

        for i, c1 in enumerate(classes):
            for j, c2 in enumerate(classes[i + 1 :], i + 1):
                debug(f"Running model for classes {c1} and {c2}")
                alpha, b, mask = models[idx]

                alpha_full = torch.zeros(
                    X_train.shape[0], device=device, dtype=alpha.dtype
                )
                alpha_full[mask] = alpha

                if K is None:
                    pred = kernel.apply(X_test, alpha_full) + b
                else:
                    pred = K @ alpha_full + b

                votes[pred < 0, i] += 1
                votes[pred >= 0, j] += 1
                idx += 1

        return classes[torch.argmax(votes, dim=1)]
