from typing import Optional

import numpy as np

from vicentin.optimization.minimization import barrier_method
from vicentin.utils.logging import info, debug


def QP(
    P: np.ndarray,
    q: np.ndarray,
    x0: np.ndarray,
    G: Optional[np.ndarray] = None,
    h: Optional[np.ndarray] = None,
    A: Optional[np.ndarray] = None,
    b: Optional[np.ndarray] = None,
    max_iter: int = 100,
    epsilon: float = 1e-4,
    mu: float = 10,
    return_dual: bool = False,
):
    n = q.size
    info(f"Starting Quadratic Programming solver [NumPy] for {n} variables.")

    obj = lambda x: 0.5 * x.dot(P @ x) + q.dot(x)
    grad_obj = lambda x: P @ x + q
    hess_obj = lambda x: P

    if A is None or b is None:
        equality = None
        n_eq = 0
    else:
        equality = (A, b)
        n_eq = A.shape[0]

    constraints = []

    if G is not None and h is not None:
        H = np.zeros((n, n))

        for i in range(len(h)):
            g_row = G[i]
            h_val = h[i]

            const = lambda x, g_row=g_row, h_val=h_val: g_row.dot(x) - h_val
            grad_const = lambda x, g_row=g_row: g_row
            hess_const = lambda x, H=H: H

            constraints.append((const, grad_const, hess_const))

    info(
        f"QP setup: {len(constraints)} inequality and {n_eq} equality constraints."
    )
    debug("Handing execution to the Barrier Method.")

    x_star, (lambdas, nu_vec) = barrier_method(
        (obj, grad_obj, hess_obj),
        constraints,
        x0,
        equality,
        max_iter,
        epsilon,
        mu,
        return_dual=True,
    )

    return (x_star, (lambdas, nu_vec)) if return_dual else x_star
