<filename>src/modules/agents/rnn_lf_agent.py
import torch
import torch.nn as nn
import torch.nn.functional as F


class RNN(nn.Module):
    """
    RNN module
    """
    def __init__(self, input_shape, args):
        super(RNN, self).__init__()
        self.args = args

        self.fc1 = nn.Linear(input_shape, args.rnn_hidden_dim)
        self.rnn = nn.GRUCell(args.rnn_hidden_dim, args.rnn_hidden_dim)

    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):
        x = F.relu(self.fc1(inputs))
        h_in = hidden_state.view(-1, self.args.rnn_hidden_dim)
        h_out = self.rnn(x, h_in)
        h_out = h_out.view(-1, self.args.n_agents, self.args.rnn_hidden_dim)
        return h_out


class FreqInfer(nn.Module):
    """
    communication frequence infer module
    """
    def __init__(self, args):
        super(FreqInfer, self).__init__()
        self.args = args
        self.fc1 = nn.Linear(args.rnn_hidden_dim, args.rnn_hidden_dim)
        self.fc2 = nn.Linear(args.rnn_hidden_dim, args.freq_infer_dim)

    def forward(self, h_out):
        x = F.relu(self.fc1(h_out))
        x = self.fc2(x)
        return x


class CommChannel(nn.Module):
    """
    communication frequence infer module
    """
    def __init__(self, args):
        super(CommChannel, self).__init__()
        self.args = args
        self.fc1 = nn.Linear(args.rnn_hidden_dim * args.n_agents, args.msg_dim)
        self.fc2 = nn.Linear(args.msg_dim, args.msg_dim)

    def forward(self, messages):
        x = F.relu(self.fc1(messages))
        x = self.fc2(x)
        return x


class QValueInfer(nn.Module):
    """
    communication frequence infer module
    """
    def __init__(self, args):
        super(QValueInfer, self).__init__()
        self.args = args
        self.comb_msg = nn.Linear(args.rnn_hidden_dim + args.msg_dim, args.rnn_hidden_dim)
        self.fc = nn.Linear(args.rnn_hidden_dim, args.n_actions)

    def forward(self, x, messages=None):
        if messages is not None:
            msg_combined = torch.cat([x, messages], dim=-1)
            x = F.relu(self.comb_msg(msg_combined))
        x = self.fc(x)
        return x


class RnnLFAgent(nn.Module):
    """
    input shape: [batch_size, in_feature]
    output shape: [batch_size, n_actions]
    hidden state shape: [batch_size, hidden_dim]
    """
    def __init__(self, input_dim, args):
        super().__init__()

        self.args = args
        self.rnn = RNN(input_dim, args)
        self.freq_infer = FreqInfer(args)
        self.comm_chennel = CommChannel(args)
        self.q_infer = QValueInfer(args)

    def q_with_communication(self, h_out, message=None):
        q_without_comm = self.q_infer(h_out, message)
        return q_without_comm

    def cuda(self):
        self.rnn.cuda()
        self.q_infer.cuda()
        self.comm_chennel.cuda()
        self.freq_infer.cuda()

    def communicate(self, hidden, comm_mask=None):
        """
        input: hidden [batch_size, n_agents, hidden_dim]
        output: messages [batch_size, n_agents, msg_dim]
        """
        if comm_mask is not None:
            comm_mask = comm_mask.float()
            comm_mask = comm_mask.unsqueeze(2).expand_as(hidden)
            hidden = hidden * comm_mask

        h_cat = hidden.view(-1, self.args.n_agents * self.args.rnn_hidden_dim)
        messages = self.comm_chennel(h_cat)
        messages = messages.unsqueeze(1).expand(-1, self.args.n_agents, -1)

        return messages
