import torch
import torch.nn as nn


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

    def forward(self, q_values: torch.Tensor, actions: torch.Tensor, rewards: torch.Tensor, next_q_values: torch.Tensor,
                dones: torch.Tensor) -> torch.Tensor:
        batch_size = q_values.shape[0]
        q_pred = q_values[torch.arange(batch_size), actions.long()]
        q_target = rewards + self.gamma * next_q_values.max(dim=1)[0] * (1 - dones)
        loss = ((q_pred - q_target) ** 2).mean()
        return loss


batch_size = 32
action_dim = 4


def get_inputs():
    q_values = torch.randn(batch_size, action_dim)
    actions = torch.randint(0, action_dim, (batch_size,))
    rewards = torch.randn(batch_size)
    next_q_values = torch.randn(batch_size, action_dim)
    dones = torch.randint(0, 2, (batch_size,)).float()
    return [q_values, actions, rewards, next_q_values, dones]


def get_init_inputs():
    gamma = torch.tensor(0.99)
    return [gamma]