"""用于 GyMRT2A 空间动作网格的 HATT actor。

动作网格中的每个格子都作为一个候选位置，并共用同一个 decoder。当前可见的其他
智能体先通过 A-A attention 汇总，再参与候选位置打分。
"""

import torch
import torch.nn as nn

from algorithms.algorithm.HATTV20Network import FeatureEncoder, RelationAttention
from algorithms.utils.distributions import FixedCategorical
from algorithms.utils.rnn import RNNLayer
from algorithms.utils.util import check


class HATTGyMRT2AActor(nn.Module):
    """将空间观测编码成网格中每个位置的动作分数。"""
    def __init__(self, args, ego_dim, other_dim, task_dim, num_agents, num_tasks, device):
        super().__init__()
        self.num_agents = num_agents
        self.num_tasks = num_tasks
        self.hidden_size = args.hidden_size
        self.use_policy_active_masks = args.use_policy_active_masks
        self.use_recurrent = args.use_naive_recurrent_policy or args.use_recurrent_policy
        self.tpdv = dict(dtype=torch.float32, device=device)
        hidden = args.hatt_hidden_dim

        self.ego_encoder = FeatureEncoder(ego_dim, hidden)
        # 当前智能体、其他智能体和候选位置的字段分别编码。
        self.other_encoder = FeatureEncoder(other_dim, hidden)
        self.candidate_encoder = FeatureEncoder(task_dim, hidden)
        self.aa_attention = RelationAttention(hidden, edge_dim=3, num_heads=args.hatt_num_heads)
        # A-A 输出与当前智能体表示合并后，再转换到 MAPPO 的 hidden_size。
        self.agent_fusion = nn.Sequential(nn.Linear(hidden * 2, hidden), nn.ReLU(), nn.Linear(hidden, hidden))
        self.agent_norm = nn.LayerNorm(hidden)
        self.agent_projection = nn.Sequential(
            nn.Linear(hidden, self.hidden_size), nn.ReLU(), nn.LayerNorm(self.hidden_size)
        )
        if self.use_recurrent:
            self.rnn = RNNLayer(
                self.hidden_size, self.hidden_size, args.recurrent_N, args.use_orthogonal
            )
        self.decoder = nn.Sequential(
            # 每个格子的输入包含智能体表示、候选位置表示和二维相对坐标。
            nn.Linear(self.hidden_size + hidden + 2, self.hidden_size),
            nn.ReLU(),
            nn.Linear(self.hidden_size, self.hidden_size),
            nn.ReLU(),
            nn.Linear(self.hidden_size, 1),
        )
        nn.init.orthogonal_(self.decoder[-1].weight, gain=0.01)
        nn.init.zeros_(self.decoder[-1].bias)
        self.last_diagnostics = {}
        self.to(device)

    def _prepare_obs(self, obs):
        """整理观测张量，并记录调用方是否只传入了一个 batch。"""
        if not isinstance(obs, dict):
            raise TypeError("HATTGyMRT2AActor requires Dict observations.")
        result = {key: check(value).to(**self.tpdv) for key, value in obs.items()}
        if result["ego"].ndim == 2:
            # buffer 会把 batch 和智能体维压平，这里恢复成两维。
            result = {
                key: value.reshape(value.shape[0] // self.num_agents, self.num_agents, *value.shape[1:])
                for key, value in result.items()
            }
        return result

    def _distribution(self, obs, rnn_states, masks, available_actions):
        # HATT 先输出任务槽分数，环境适配层已将任务槽排成空间动作网格。
        # RNN 保存跨时间步的信息，动作掩码负责去掉当前不可达的位置。
        obs = self._prepare_obs(obs)
        batch, agents, _ = obs["ego"].shape
        available = check(available_actions).to(**self.tpdv).reshape(batch, agents, self.num_tasks)

        ego_h = self.ego_encoder(obs["ego"])
        other_h = self.other_encoder(obs["others"])
        candidate_h = self.candidate_encoder(obs["tasks"])
        # others 的前两项是相对位置，第 3 项表示该智能体是否可见。
        relative = obs["others"][..., :2]
        visible = obs["others"][..., 2] > 0.5
        edge = torch.cat([relative, visible.float().unsqueeze(-1)], dim=-1)
        aa_context, aa_weights = self.aa_attention(ego_h, other_h, edge, visible)
        # 残差保留当前智能体自己的编码。
        agent_h = self.agent_norm(ego_h + self.agent_fusion(torch.cat([ego_h, aa_context], dim=-1)))
        agent_h = self.agent_projection(agent_h).reshape(batch * agents, self.hidden_size)

        rnn_states = check(rnn_states).to(**self.tpdv)
        masks = check(masks).to(**self.tpdv)
        if self.use_recurrent:
            # RNN 的输入按 batch × 智能体展开。
            agent_h, rnn_states = self.rnn(agent_h, rnn_states, masks)
        agent_h = agent_h.reshape(batch, agents, self.hidden_size)

        expanded = agent_h.unsqueeze(2).expand(-1, -1, self.num_tasks, -1)
        # 将当前智能体表示复制到每个网格位置，再逐格计算 logit。
        candidate_xy = obs["tasks"][..., :2]
        logits = self.decoder(torch.cat([expanded, candidate_h, candidate_xy], dim=-1)).squeeze(-1)
        logits = logits.reshape(batch * agents, self.num_tasks)
        logits = logits.masked_fill(available.reshape(batch * agents, self.num_tasks) <= 0, -1e10)
        # 诊断项只用于日志，全部从反向传播图中分离。
        self.last_diagnostics = {
            "aa_attention": aa_weights.detach(),
            "visible_agents": visible.float().sum(dim=-1).detach(),
            "candidate_context_norm": candidate_h.norm(dim=-1).detach(),
        }
        return FixedCategorical(logits=logits), rnn_states

    def forward(self, obs, rnn_states, masks, available_actions=None, deterministic=False):
        dist, rnn_states = self._distribution(obs, rnn_states, masks, available_actions)
        actions = dist.mode() if deterministic else dist.sample()
        return actions, dist.log_probs(actions), rnn_states

    def evaluate_actions(
        self,
        obs,
        rnn_states,
        action,
        masks,
        available_actions=None,
        active_masks=None,
        gate_actions=None,
    ):
        del gate_actions  # GyMRT2A 只有一个离散动作输出，不使用额外 gate 动作。
        dist, _ = self._distribution(obs, rnn_states, masks, available_actions)
        action = check(action).to(device=self.tpdv["device"])
        log_probs = dist.log_probs(action)
        entropy = dist.entropy()
        # 未激活的智能体不参与平均熵计算。
        if active_masks is not None and self.use_policy_active_masks:
            active_masks = check(active_masks).to(**self.tpdv).squeeze(-1)
            entropy = (entropy * active_masks).sum() / active_masks.sum().clamp_min(1.0)
        else:
            entropy = entropy.mean()
        return log_probs, entropy