from typing import Optional
import torch

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


def LP(
    c: torch.Tensor,
    x0: torch.Tensor,
    G: Optional[torch.Tensor] = None,
    h: Optional[torch.Tensor] = None,
    A: Optional[torch.Tensor] = None,
    b: Optional[torch.Tensor] = None,
    max_iter: int = 100,
    epsilon: float = 1e-4,
    mu: float = 10,
    return_dual: bool = False,
):
    info(
        f"Starting Linear Programming solver [Torch] with {c.numel()} variables on {c.device}."
    )

    obj = lambda x: c.dot(x)

    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:
        for i in range(len(h)):
            c_func = lambda x, g_row=G[i], h_val=h[i]: g_row.dot(x) - h_val
            constraints.append(c_func)

    info(
        f"LP setup: {len(constraints)} inequality and {n_eq} equality constraints."
    )
    debug("Passing control to Barrier Method.")

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

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