import torch


def fista_momentum(cur_Z, prev_Z, momentum):
    """
    Calculates a linear combination of the last two sparse codings with a momentum term

    :param cur_Z: Sparse code found in current step
    :param prev_Z: Sparse code found in previous step
    :param momentum: float. Momentum term.
    :return: Updated sparse code to be used in next step.
    """
    next_momentum = (1 + torch.sqrt(1+4*(momentum**2)))/2
    momentum_ratio = (momentum - 1) / next_momentum

    pushed_Z = (cur_Z - prev_Z) * momentum_ratio
    next_Z = cur_Z + pushed_Z

    return next_Z, next_momentum