from typing import Dict
from typing import Counter


def calc_nstep_return(
    r_t: Dict[str, float], discounts: list, rewards: list
) -> Dict[str, float]:
    """Function that calculates n_step_return as follows:
    R_{t:t+n} := r_t + d_t * r_{t+1} + ...    + d_t * ... * d_{t+n-2} * r_{t+n-1}.

    Args:
        r_t (Dict[str,float]): reward achieved from action a_t.
        discounts (list): list of discounts.
        rewards (list): list of rewards.

    Returns:
        A [Dict[str,float] with the return per agent.
    """
    reward = Counter(r_t)
    for index, r_tn in enumerate(rewards):
        d_t = discounts[index]
        return_t = {key: d_t[key] * r_tn.get(key, 0) for key in r_tn.keys()}
        reward.update(Counter(return_t))
    return dict(reward)