import torch


def get_gae(value, reward, args, is_normalize=False, eps=1e-8):
    """Compute generalized advantage estimator

    Args:
        value (list): Contains value function across trajectories
        reward (list): Contains rewards across trajectories for specific agent
        args (argparse): Python argparse that contains arguments
        is_normalize (bool): Normalize baseline if flag is True. Default: False
        eps (float): Epsilon for numerical stability. Default: 1e-8

    Returns:
        GAE (torch.Tensor): Estimated generalized advantage function

    References:
        https://github.com/dgriff777/rl_a3c_pytorch/blob/master/train.py
    """
    value = torch.stack(value, dim=1)
    assert value.shape == (args.traj_batch_size, args.ep_horizon), \
        "Shape must be: (batch, ep_horizon)"
    value = torch.cat((value, torch.zeros(value.shape[0], 1)), dim=1)
    reward = torch.stack(reward, dim=1)
    assert reward.shape == (args.traj_batch_size, args.ep_horizon), \
        "Shape must be: (batch, ep_horizon)"

    gae, advantage = 0., []
    for timestep in reversed(range(args.ep_horizon)):
        delta = (reward[:, timestep] + args.discount * value[:, timestep + 1]) - value[:, timestep]
        gae = gae * args.discount * args.lambda_ + delta
        advantage.insert(0, gae)
    advantage = torch.stack(advantage, dim=1)
    assert reward.shape == advantage.shape

    if is_normalize:
        advantage = advantage - torch.mean(advantage)
        std = torch.sqrt(torch.mean(advantage ** 2))
        advantage.div_(std + eps)

    return advantage