"""Monotonic mixing network used by the local QMIX implementation."""

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


class QMixer(nn.Module):
    """Map per-agent utilities and the centralized state to a joint Q value.

    Absolute hypernetwork weights enforce the monotonicity constraint from
    QMIX, so decentralized greedy actions maximize the centralized value.
    Inputs use ``[..., num_agents]`` and ``[..., state_dim]`` shapes.
    """

    def __init__(self, state_dim, num_agents, hidden_dim=32, hypernet_hidden_dim=64):
        super().__init__()
        self.state_dim = int(state_dim)
        self.num_agents = int(num_agents)
        self.hidden_dim = int(hidden_dim)
        hypernet_hidden_dim = int(hypernet_hidden_dim)
        self.hyper_w1 = nn.Sequential(
            nn.Linear(self.state_dim, hypernet_hidden_dim),
            nn.ReLU(),
            nn.Linear(hypernet_hidden_dim, self.num_agents * self.hidden_dim),
        )
        self.hyper_b1 = nn.Linear(self.state_dim, self.hidden_dim)
        self.hyper_w2 = nn.Sequential(
            nn.Linear(self.state_dim, hypernet_hidden_dim),
            nn.ReLU(),
            nn.Linear(hypernet_hidden_dim, self.hidden_dim),
        )
        self.hyper_b2 = nn.Sequential(
            nn.Linear(self.state_dim, hypernet_hidden_dim),
            nn.ReLU(),
            nn.Linear(hypernet_hidden_dim, 1),
        )

    def forward(self, agent_qs, states):
        prefix = agent_qs.shape[:-1]
        flat_qs = agent_qs.reshape(-1, 1, self.num_agents)
        flat_states = states.reshape(-1, self.state_dim)
        w1 = self.hyper_w1(flat_states).abs().view(
            -1, self.num_agents, self.hidden_dim
        )
        b1 = self.hyper_b1(flat_states).view(-1, 1, self.hidden_dim)
        hidden = F.elu(torch.bmm(flat_qs, w1) + b1)
        w2 = self.hyper_w2(flat_states).abs().view(-1, self.hidden_dim, 1)
        b2 = self.hyper_b2(flat_states).view(-1, 1, 1)
        return (torch.bmm(hidden, w2) + b2).view(*prefix, 1)