import torch


def ppostep(policy, optimizer, states, actions, baseprobs, values, advantages, epscut, gradclip,
            valuecoef, entcoef):
    """Performs a step of Proximal Policy Optimization

    Arguments:
        - policy: policy network to optimize.
        - optimizer: pytorch optimizer algorithm to use.
        - states: iterable of gathered experience states
        - actions: iterable of performed actions in the states
        - basepros: current base probabilities of performing those actions
        - values: estimated values for each state
        - advantages: estimated advantage values for those actions
        - epscut: epsilon cut for policy gradient update
        - gradclip: maximum norm for clipping gradients
        - valuecoef: weight of value function loss
        - entcoef: weight of entropy function loss

    Returns the value of the losses computed in the optimization step.

    Reference: https://arxiv.org/pdf/1707.06347.pdf
    """
    optimizer.zero_grad()
    # Compute action probabilities and state values for current network parameters
    logprobs, newvalues = zip(*[policy.action_logprobs_value(st, ac) for st, ac in zip(states, actions)])
    newvalues = torch.stack([x[0][0] for x in newvalues])
    # Policy Gradients loss (advantages)
    newprobs = [torch.exp(logp) for logp in logprobs]
    probratios = [newprob / prob for prob, newprob in zip(baseprobs, newprobs)]
    clippings = [torch.min(rt * adv, torch.clamp(rt, 1 - epscut, 1 + epscut) * adv)
                 for rt, adv in zip(probratios, advantages)]
    pgloss = -torch.cat(clippings).mean()
    # Entropy loss
    entropyloss = - entcoef * torch.mean(torch.stack([policy.entropy(st) for st in states]))
    # Value estimation loss
    # (newvalue - [advantage + oldvalue])^2, that is, make new value closer to estimated error in old estimate
    # A clipped version of the loss is also included
    returns = advantages + values
    valuelosses = (newvalues - returns) ** 2
    newvalues_clipped = values + torch.clamp(newvalues - values, epscut)
    valuelosses_clipped = (newvalues_clipped - returns) ** 2
    valueloss = valuecoef * 0.5 * torch.mean(torch.max(valuelosses, valuelosses_clipped))
    # Total loss
    loss = pgloss + valueloss + entropyloss
    # Optimizer step, with clipped gradients to a max norm to avoid exploding gradients
    loss.backward()
    torch.nn.utils.clip_grad_norm(policy.parameters(), gradclip)
    optimizer.step()

    return loss, pgloss, valueloss, entropyloss