import torch.nn as nn
import torch.nn.functional as F
import torch as th

class RNNAgent(nn.Module):
    def __init__(self, input_shape, args):
        super(RNNAgent, self).__init__()
        self.args = args

        self.fc1 = nn.Linear(input_shape, args.rnn_hidden_dim)
        if self.args.use_rnn:
            self.rnn = nn.GRUCell(args.rnn_hidden_dim, args.rnn_hidden_dim)
        self.fc2 = nn.Linear(args.rnn_hidden_dim, args.n_actions)
        self.fc3 = nn.Linear(args.rnn_hidden_dim, args.n_actions)

    def init_hidden(self):
        # make hidden states on same device as model
        return self.fc1.weight.new(1, self.args.rnn_hidden_dim).zero_()

    def forward(self, inputs, hidden_state):
        
        if self.args.use_rnn:
            x = F.relu(self.fc1(inputs))
            h_in = hidden_state.reshape(-1, self.args.rnn_hidden_dim)
            h = self.rnn(x, h_in)
            q1 = self.fc2(h)
            q2 = self.fc3(h)
            return th.min(q1, q2), hidden_state
        else:
            x = F.relu(self.fc1(inputs))
            q1 = self.fc2(x)
            q2 = self.fc3(x)
            return th.min(q1, q2), hidden_state


    def init_weights(self, m):
        if type(m) == nn.Linear:
            nn.init.uniform_(m.weight)
            m.bias.data.fill_(0.0)
            
