import torch


def retrace(q_values, rewards, importance_weights=None, discount=0.99, l=0.75):
    """
    Retrace estimate (Munos et al., 2016).

    Args:
        q_values (torch.Tensor): the Q-value estimates at each time step [time_steps+1, batch_size, 1]
        rewards (torch.Tensor): the rewards at each time step [time_steps, batch_size, 1]
        importance_weights (torch.Tensor): the importance weights at each time step [time_steps, batch_size, 1]
        discount (float): the temporal discount factor
        l (float): the lambda weighting factor
    """
    if q_values.shape[0] == 1 or l == -1:
        # degenerate case
        return q_values

    if importance_weights is None:
        # On-policy
        importance_weights = torch.ones_like(q_values)

    deltas = rewards + discount * q_values[1:] - q_values[:-1]
    importance_weights = l * torch.clamp(importance_weights, 0, 1)[:-2]
    importance_weights = torch.cat([torch.ones_like(q_values[:1]), importance_weights], 0)
    discounts = torch.cat([(discount*torch.ones_like(q_values[:1]))**i for i in range(q_values.shape[0])], 0)
    q_estimates = q_values[:1] + torch.sum(discounts[:-1] * torch.cumprod(importance_weights, 0) * deltas, 0, keepdim=True)

    return q_estimates