import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self, max_kl):
        super(Model, self).__init__()
        self.max_kl = max_kl

    def forward(self, log_probs: torch.Tensor, old_log_probs: torch.Tensor, advantages: torch.Tensor,
                old_probs: torch.Tensor, new_probs: torch.Tensor) -> torch.Tensor:
        ratio = torch.exp(log_probs - old_log_probs)
        surrogate_loss = -(ratio * advantages).mean()
        kl_div = (old_probs * (torch.log(old_probs) - torch.log(new_probs))).sum(dim=-1).mean()
        loss = surrogate_loss + self.max_kl * kl_div
        return loss


batch_size = 32
action_dim = 4


def get_inputs():
    log_probs = torch.randn(batch_size)
    old_log_probs = torch.randn(batch_size)
    advantages = torch.randn(batch_size)
    old_probs = torch.softmax(torch.randn(batch_size, action_dim), dim=-1)
    new_probs = torch.softmax(torch.randn(batch_size, action_dim), dim=-1)
    return [log_probs, old_log_probs, advantages, old_probs, new_probs]


def get_init_inputs():
    max_kl = torch.tensor(0.01)
    return [max_kl]