from typing import Sequence, Optional

import torch

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


def SOCP(
    f: torch.Tensor,
    socp_constraints: Sequence[Sequence[torch.Tensor]],
    x0: torch.Tensor,
    F: Optional[torch.Tensor] = None,
    g: Optional[torch.Tensor] = None,
    max_iter: int = 100,
    epsilon: float = 1e-4,
    mu: float = 6,
    return_dual: bool = False,
):
    info(
        f"Starting SOCP solver [Torch] with {f.numel()} variables on {f.device}."
    )

    obj = lambda x: f.dot(x)

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

    G = []
    for i, (A, b, c, d) in enumerate(socp_constraints):
        constraint = (
            lambda x, A=A, b=b, c=c, d=d: torch.sum((A @ x + b) ** 2)
            - (c.dot(x) + d) ** 2
        )

        upper_cone = lambda x, c=c, d=d: -c.dot(x) - d

        G.append(constraint)
        G.append(upper_cone)
        debug(
            f"Added SOCP constraint {i}: A shape {list(A.shape)}, c shape {list(c.shape)}"
        )

    info(
        f"SOCP setup: {len(socp_constraints)} cone constraints and {n_eq} equalities."
    )
    debug("Passing control to Barrier Method.")

    x_star, (lambdas, mu_vec) = barrier_method(
        obj,
        G,
        x0,
        equality,
        max_iter,
        epsilon,
        mu,
        return_dual=True,
    )

    u = []
    v = []

    for i in range(len(socp_constraints)):
        A, b, c, d = socp_constraints[i]
        l1 = lambdas[2 * i]  # Multiplier for quadratic form
        l2 = lambdas[2 * i + 1]  # Multiplier for (c^\top x + d) >= 0

        u_i = -2 * l1 * (A @ x_star + b)
        v_i = 2 * l1 * (c.dot(x_star) + d) + l2

        u.append(u_i)
        v.append(v_i)

    return (x_star, (u, v, mu_vec)) if return_dual else x_star
