import torch


def woodbury_solve(rhs, scaled_inv_diag_umat, woodbury_factor, scaled_inv_diag, inv_scale):
    """
    Solves the system of equations: :math:`(D + U V^T)x = b` using the Woodbury formula,
    where x is the right-hand-side (size n), U, V are (n x k), and D is a (n x n) diagonal matrix.
    Can also be used in batch mode, where U, V, and D are batches of matrices and rhs is a batch of right-hand-side.

    This should be used after calling woodbury_factor.

    Args:
        :attr:`rhs` (size n x t)
            Right hand side vector b to solve with.
        :attr:`scaled_inv_diag_umat` (n x k)
            The `D^{-1} U` matrix
        :attr:`woodbury_factor` (n x k)
            The result of calling woodbury_factor on U, V, and D.
        :attr:`diag` (vector)
            The diagonal of D
    """
    # (D + UV^T)x = D^-1 x - D^-1 U ((I + V^T D^-1 U)^-1 V^T) D^-1 x
    #             = D^-1 x - D^-1 U (1/s (woodbury_factor)) D^-1 x
    #             = s( E^-1 x - E^-1 U (woodbury_factor) E^-1 x )
    scaled_inv_diag_rhs = rhs * scaled_inv_diag
    res = torch.add(scaled_inv_diag_rhs, -1, scaled_inv_diag_umat @ (woodbury_factor @ scaled_inv_diag_rhs))
    res = res.div_(inv_scale)

    # Reshape the result to be the correct shape
    return res