import torch


def condat_vu_iteration(iterate_p, iterate_d, y, operator,
                        prior, sigma, tau, lambd, inpainting=False):
    """
    Condat Vu iteration

    Parameters
    ----------
    iterate_p : torch.Tensor
        Current primal iterate
    iterate_d : torch.Tensor
        Current dual iterate
    y : torch.Tensor
        Data
    operator : torch.Tensor
        Measurement matrix
    prior : torch.Tensor
        Prior
    sigma : torch.Tensor, shape (1)
        Step size
    tau : torch.Tensor, shape (1)
        Step size
    lambd : float
        Regularization parameter

    Returns
    -------
    out_p : torch.Tensor
        Next primal iterate
    out_d : torch.Tensor
        Next dual iterate
    """
    # Gradient descent primal
    if inpainting:
        first_gradient = operator * (iterate_p - y)
    else:
        first_gradient = torch.matmul(
            operator.t(), torch.matmul(operator, iterate_p) - y
            )
    second_gradient = torch.matmul(prior.t(), iterate_d)
    out_p = iterate_p - tau * (first_gradient + second_gradient)

    # Gradient ascent dual
    gradient_dual = torch.matmul(prior, 2 * out_p - iterate_p)

    ascent = iterate_d + sigma * gradient_dual
    out_d = ascent / torch.max(torch.ones_like(ascent),
                               torch.abs(ascent) / lambd)
    return out_p, out_d