from typing import Optional
from itertools import combinations

import numpy as np
from cvxopt import matrix, solvers

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):
    pos_mask = y == 1
    neg_mask = y == -1
    n_pos = pos_mask.sum()
    n_neg = neg_mask.sum()

    epsilon = 0.01 if np.isinf(C) else 0.05 * C

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

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

    if not np.isinf(C) and val_neg >= C:
        scale = (0.9 * C) / val_neg
        alpha0 *= scale

    return alpha0


def _train_binary(
    X,
    y,
    kernel,
    lambda_,
    C,
    smo_threshold,
    vicentin_qp,
    **kwargs,
):
    n = int(X.shape[0])
    if C is None:
        C = 1 / (lambda_ * n)

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

    debug(f"Training binary classifier with QP (n={n}, C={C:.2e})")
    y = np.asarray(y, dtype=float)

    P = kernel[:, :] if isinstance(kernel, np.ndarray) else kernel(X, X)
    q = -y

    diag_y = np.diag(y)
    if np.isinf(C):
        G = -diag_y
        h = np.zeros(n)
    else:
        G = np.vstack((-diag_y, diag_y))
        h = np.hstack((np.zeros(n), np.ones(n) * C))

    A = np.ones((1, n))
    b_eq = np.zeros(1)

    if vicentin_qp:
        debug("Using vicentin_qp solver")
        alpha0 = _get_alpha0(y, C)
        alpha_y, duals = QP(
            P, q, alpha0, G, h, A, b_eq, return_dual=True, **kwargs
        )
        b_aux = duals[1]
    else:
        debug("Using CVXOPT solver")
        solvers.options["show_progress"] = False

        P = matrix(P)
        q = matrix(q)
        G = matrix(G)
        h = matrix(h)
        A = matrix(A, (1, n))
        b_eq = matrix(0.0)

        sol = solvers.qp(P, q, G, h, A, b_eq)

        alpha_y = np.array(sol["x"]).flatten()
        b_aux = sol["y"]

    alpha = alpha_y * y
    sv_mask = (alpha > 1e-4) & (alpha < C - 1e-4)
    debug(f"Support Vectors found: {np.sum(sv_mask)}")

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

    return alpha_y, b


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

    y = y.flatten()
    classes = np.unique(y).astype(np.int32)
    info(f"Classes detected: {classes.tolist()}")

    if len(classes) == 2:
        y = np.where(y == classes[0], -1, 1)

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

        mask = np.ones_like(y, dtype=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, 2):
            debug(f"Training OvO pair: {c1} vs {c2}")
            mask = (y == c1) | (y == c2)

            y_bin = np.where(y[mask] == c1, -1, 1)
            sub_kernel = SubsetKernel(kernel, mask)

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

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

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

        for c in classes:
            debug(f"Training OvR for class: {c}")
            mask = np.ones_like(y, dtype=bool)
            y_bin = np.where(y == c, 1, -1)

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

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

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


def predict(
    model_data,
    X_train: np.ndarray | Kernel,
    X_test: np.ndarray,
    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 = kernel.X
    else:
        kernel = LinearKernel(X_train)

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

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

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

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

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

        return classes[np.argmax(scores, axis=1)]

    if strategy == "ovo":
        info("Prediction using OvO strategy")
        votes = np.zeros((X_test.shape[0], len(classes)))
        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 = np.zeros(X_train.shape[0], 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[np.argmax(votes, axis=1)]
