"""Compatibility wrapper around the maintained recurrent individual-Q net."""

import torch

from algorithms.qmix.algorithm.QMixPolicy import RecurrentAgentQNetwork


class AgentQFunction(RecurrentAgentQNetwork):
    def __init__(self, args, input_dim, act_dim, device):
        hidden_dim = int(
            getattr(args, "qmix_agent_hidden_dim", getattr(args, "hidden_size", 128))
        )
        super().__init__(input_dim, act_dim, hidden_dim)
        self.device = device
        self.to(device)

    def forward(self, obs, rnn_states):
        obs = torch.as_tensor(obs, dtype=torch.float32, device=self.device)
        hidden = torch.as_tensor(rnn_states, dtype=torch.float32, device=self.device)
        if obs.ndim == 2:
            if hidden.ndim == 3:
                hidden = hidden[0]
            return self.forward_step(obs, hidden)
        if hidden.ndim == 3:
            hidden = hidden[0]
        outputs = []
        for step in range(obs.shape[0]):
            q_values, hidden = self.forward_step(obs[step], hidden)
            outputs.append(q_values)
        return torch.stack(outputs), hidden