from itertools import combinations

import numpy as np

from vicentin.kernels import SubsetKernel
from vicentin.machine_learning.regression import RidgeRegression
from vicentin.utils.logging import debug, info


def _sigmoid(u):
    return 1 / (1 + np.exp(-u))


def _train_binary(X_kernel, y, lambda_orig, max_iter, tol, is_primal):
    if is_primal:
        N = X_kernel.shape[0]
        dtype = X_kernel.dtype
    else:
        N = X_kernel.N
        dtype = X_kernel.X.dtype

    m = np.zeros((N, 1), dtype=dtype)
    y = y.reshape(-1, 1) if y.ndim == 1 else y

    weights = None
    ridge_model = RidgeRegression(lambda_=lambda_orig, backend="numpy")

    for i in range(max_iter):
        m_prev = m.copy()

        W = _sigmoid(m) * _sigmoid(-m)
        z = m + y * (1 + np.exp(-y * m))

        ridge_model.fit(X_kernel, z, W=W)

        X_pred = X_kernel if is_primal else X_kernel.X
        m = ridge_model.predict(X_pred).reshape(-1, 1)

        diff = float(np.max(np.abs(m - m_prev)))
        debug(f"Logistic Iteration {i+1}/{max_iter}, Diff: {diff:.6f}")

        if diff < tol:
            info(f"Logistic Regression converged in {i+1} iterations.")
            break

    weights = ridge_model.weights

    if is_primal and not ridge_model.is_primal:
        weights = X_kernel.T @ weights

    return weights


def logistic_regression(
    X_kernel, y, lambda_orig, max_iter, tol, is_primal, strategy="ovo"
):
    y = y.flatten()
    classes = np.unique(y).astype(np.int32)
    info(f"Classes detected: {classes.tolist()}")

    X = X_kernel if is_primal else X_kernel.X

    if len(classes) == 2:
        y_bin = np.where(y == classes[0], -1, 1)
        weights = _train_binary(
            X_kernel, y_bin, lambda_orig, max_iter, tol, is_primal
        )
        mask = np.ones_like(y, dtype=bool)
        return {
            "classes": classes,
            "models": [(weights, 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)

            if is_primal:
                X_sub = X_kernel[mask]
            else:
                X_sub = SubsetKernel(X_kernel, mask)

            weights = _train_binary(
                X_sub, y_bin, lambda_orig, max_iter, tol, is_primal
            )
            models.append((weights, 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)
            weights = _train_binary(
                X_kernel, y_bin, lambda_orig, max_iter, tol, is_primal
            )
            models.append((weights, mask))

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


def predict(model_data, kernel, X_test, is_primal):
    classes = model_data["classes"]
    models = model_data["models"]
    strategy = model_data["strategy"]

    def get_pred(weights, mask):
        if is_primal:
            return X_test @ weights
        else:
            weights_full = np.zeros((kernel.N, 1), dtype=weights.dtype)
            weights_full[mask] = weights
            return kernel.apply(X_test, weights_full)

    if strategy == "binary":
        info("Prediction using binary strategy")
        pred = get_pred(models[0][0], models[0][1]).flatten()
        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, (weights, mask) in enumerate(models):
            debug(f"Running model {i}")
            scores[:, i] = get_pred(weights, mask).flatten()

        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

        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}")

                weights, mask = models[idx]
                pred = get_pred(weights, mask).flatten()
                votes[pred < 0, i] += 1
                votes[pred >= 0, j] += 1
                idx += 1

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