import numpy as np

from vicentin.optimization.minimization import conjugate_gradient
from vicentin.utils.logging import debug


def ridge_regression(X_or_kernel, y, W, lambda_reg, is_primal):
    if is_primal:
        X = X_or_kernel
        N, d = X.shape

        y = y.reshape(-1, 1) if y.ndim == 1 else y
        W = (
            np.ones((N, 1), dtype=X.dtype)
            if W is None
            else W.reshape(-1, 1) if W.ndim == 1 else W
        )

        b = X.T @ (W * y)

        if d < 10_000:
            debug(f"Solving dense Primal Ridge for d={d}")
            A = X.T @ (W * X) + lambda_reg * np.eye(d, dtype=X.dtype)
            weights = np.linalg.solve(A, b)
        else:
            debug(f"Solving CG Primal Ridge for d={d}")
            A_op = lambda v: (X.T @ (W * (X @ v))) + lambda_reg * v
            weights = conjugate_gradient(A_op, b, max_iter=d)

        return weights

    else:
        kernel = X_or_kernel
        N = kernel.N

        y = y.reshape(-1, 1) if y.ndim == 1 else y
        W = (
            np.ones((N, 1), dtype=y.dtype)
            if W is None
            else W.reshape(-1, 1) if W.ndim == 1 else W
        )

        if N < 10_000:
            debug(f"Solving dense Dual Ridge for N={N}")
            K = kernel[:, :]
            A = (W * K) + lambda_reg * np.eye(N, dtype=y.dtype)
            alpha = np.linalg.solve(A, W * y)
        else:
            debug(f"Solving CG Dual Ridge for N={N}")
            A_op = lambda v: W * (kernel @ v) + lambda_reg * v
            alpha = conjugate_gradient(A_op, W * y, max_iter=N)

        return alpha
