"""
模型组件：TTA头、策略头
"""

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

class TTAHead(nn.Module):
    """
    TTA回归头
    输入: belief向量 [B, hidden_dim]
    输出: (tta_mean, tta_logvar)
    """
    def __init__(self, hidden_dim, intermediate_dim=512):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.intermediate_dim = intermediate_dim
        
        self.net = nn.Sequential(
            nn.Linear(hidden_dim, intermediate_dim),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(intermediate_dim, 128),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(128, 2)  # mean, log_var
        )
    
    def forward(self, hidden_state):
        """
        Args:
            hidden_state: [B, hidden_dim]
        Returns:
            tta_mean: [B]
            tta_logvar: [B]
        """
        output = self.net(hidden_state)
        tta_mean = output[:, 0]
        tta_logvar = output[:, 1]
        return tta_mean, tta_logvar


class PolicyHead(nn.Module):
    """
    策略头(DPO阶段训练)
    输入: belief向量 + TTA统计 + 历史编码
    输出: 动作logits [B, 3]
    """
    def __init__(self, hidden_dim, num_actions=3, dropout=0.2):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.num_actions = num_actions

        # 历史动作编码器
        self.action_embedding = nn.Embedding(num_actions, 16)

        # 策略网络
        # 输入: hidden_dim + 2(tta_mean, tta_var) + 16(history)
        input_dim = hidden_dim + 2 + 16

        self.net = nn.Sequential(
            nn.Linear(input_dim, 512),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(256, num_actions)
        )
    
    def forward(self, hidden_state, tta_mean, tta_var, prev_action):
        """
        Args:
            hidden_state: [B, hidden_dim]
            tta_mean: [B]
            tta_var: [B]
            prev_action: [B] (0=silent, 1=observe, 2=alert)
        Returns:
            action_logits: [B, 3]
        """
        # 编码历史动作
        action_emb = self.action_embedding(prev_action)  # [B, 16]
        
        # 拼接所有特征
        features = torch.cat([
            hidden_state,
            tta_mean.unsqueeze(-1),
            tta_var.unsqueeze(-1),
            action_emb
        ], dim=-1)
        
        logits = self.net(features)
        return logits


class EvidentialPolicyHead(nn.Module):
    """
    Evidential PolicyHead — outputs Dirichlet concentration parameters α.

    Instead of softmax logits, predicts evidence e ≥ 0 for each class,
    then α = e + 1 forms a Dirichlet distribution Dir(α).

    From α we derive:
      - expected probability:  p = α / S  where S = Σα
      - epistemic uncertainty: u = K / S  (K = num_actions)

    At inference, high u → default to OBSERVE (conservative).
    """

    def __init__(self, hidden_dim, num_actions=3, dropout=0.2):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.num_actions = num_actions

        self.action_embedding = nn.Embedding(num_actions, 16)

        input_dim = hidden_dim + 2 + 16

        self.net = nn.Sequential(
            nn.Linear(input_dim, 512),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(512, 256),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(256, num_actions),
        )

        nn.init.zeros_(self.net[-1].weight)
        nn.init.constant_(self.net[-1].bias, 1.0)

    def forward(self, hidden_state, tta_mean, tta_var, prev_action):
        action_emb = self.action_embedding(prev_action)
        features = torch.cat([
            hidden_state,
            tta_mean.unsqueeze(-1),
            tta_var.unsqueeze(-1),
            action_emb,
        ], dim=-1)
        out = self.net(features)
        evidence = F.softplus(out)
        alpha = evidence + 1.0
        return alpha

    def predict(self, alpha):
        S = alpha.sum(dim=-1, keepdim=True)
        p = alpha / S
        u = float(self.num_actions) / S.squeeze(-1)
        return p, u


class BinaryCollisionHead(nn.Module):
    """Binary collision classifier for Nexar-style detection.
    Bypasses 3-class softmax bottleneck by directly predicting P(collision)."""

    def __init__(self, hidden_dim=2048, dropout=0.2):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(hidden_dim + 2, 512),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(512, 256),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(256, 1),
        )

    def forward(self, hidden_state, tta_mean, tta_var):
        x = torch.cat([hidden_state,
                        tta_mean.unsqueeze(-1),
                        tta_var.unsqueeze(-1)], dim=-1)
        return self.net(x).squeeze(-1)


class BinaryTemporalHead(nn.Module):
    """Per-window binary collision scorer with max aggregation (BADAS-style)."""

    def __init__(self, hidden_dim=2048, proj_dim=256, dropout=0.2):
        super().__init__()
        self.proj = nn.Linear(hidden_dim, proj_dim)
        self.scorer = nn.Sequential(
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(proj_dim + 2, 128),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(128, 1),
        )

    def forward(self, beliefs_frame, tta_mean_seq=None, tta_var_seq=None,
                valid_mask=None):
        """
        beliefs_frame: [B, T, D]
        tta_mean_seq:  [B, T] or None
        tta_var_seq:   [B, T] or None
        valid_mask:    [B, T] or None
        Returns: clip_score [B], per_window_score [B, T]
        """
        B, T, D = beliefs_frame.shape
        h = self.proj(beliefs_frame)
        if tta_mean_seq is not None:
            h = torch.cat([h,
                           tta_mean_seq.unsqueeze(-1),
                           tta_var_seq.unsqueeze(-1)], dim=-1)
        else:
            h = torch.cat([h, torch.zeros(B, T, 2, device=h.device)], dim=-1)
        per_window = self.scorer(h).squeeze(-1)
        if valid_mask is not None:
            per_window = per_window.masked_fill(~valid_mask, -1e9)
        clip_score = per_window.max(dim=1).values
        return clip_score, per_window


class HierarchicalPolicyHead(nn.Module):
    """
    Hierarchical Risk Assessment Head — replaces 3-class softmax with two
    independent binary classifiers to break probability competition.

    Motivation (empirical + theoretical):
      - 3-class softmax locks AP at 0.24 because P(ALERT) + P(OBSERVE) + P(SILENT) = 1,
        so high P(OBSERVE) necessarily suppresses P(ALERT).
      - Binary ablation (OBSERVE→ALERT merge) achieves AP=0.888, proving features
        are sufficient — the bottleneck is the output parameterisation.
      - Binary Relevance decomposition (Tsoumakas & Katakis, 2007; Read et al., 2011)
        avoids label competition inherent in shared-simplex classifiers.
      - Hierarchical decision-making aligns with cascaded safety assessment in AD
        (Norden et al., 2025; Pjetri et al., ECCV-W 2025).

    Architecture:
      SharedTrunk: (belief ⊕ tta_mean ⊕ tta_var ⊕ action_emb) → 512 → 256
      AlertHead:  256 → 1 (sigmoid)  — P(ALERT)  — "immediate danger"
      DangerHead: 256 → 1 (sigmoid)  — P(DANGER) — "any non-SILENT response needed"

    Decision logic:
      P(ALERT) > τ_a  → ALERT
      P(DANGER) > τ_d → OBSERVE
      else             → SILENT
    """

    def __init__(self, hidden_dim, dropout=0.2):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.action_embedding = nn.Embedding(3, 16)

        input_dim = hidden_dim + 2 + 16  # belief + tta_mean + tta_var + action_emb

        self.shared = nn.Sequential(
            nn.Linear(input_dim, 512),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(512, 256),
            nn.GELU(),
            nn.Dropout(dropout),
        )

        # Independent binary outputs (logit space — apply sigmoid externally)
        self.alert_head = nn.Linear(256, 1)
        self.danger_head = nn.Linear(256, 1)

        # Balanced init
        nn.init.zeros_(self.alert_head.weight)
        nn.init.zeros_(self.alert_head.bias)
        nn.init.zeros_(self.danger_head.weight)
        nn.init.zeros_(self.danger_head.bias)

    def forward(self, hidden_state, tta_mean, tta_var, prev_action):
        """
        Returns:
            alert_logit:  [B] — raw logit for ALERT
            danger_logit: [B] — raw logit for DANGER (OBSERVE+ALERT vs SILENT)
        """
        action_emb = self.action_embedding(prev_action)
        features = torch.cat([
            hidden_state,
            tta_mean.unsqueeze(-1),
            tta_var.unsqueeze(-1),
            action_emb,
        ], dim=-1)
        h = self.shared(features)
        alert_logit = self.alert_head(h).squeeze(-1)    # [B]
        danger_logit = self.danger_head(h).squeeze(-1)   # [B]
        return alert_logit, danger_logit

    def predict(self, alert_logit, danger_logit, tau_alert=0.5, tau_danger=0.5):
        """
        Hierarchical decision with configurable thresholds.
        Returns:
            preds:    [B] long  — 0=SILENT, 1=OBSERVE, 2=ALERT
            p_alert:  [B] float — sigmoid probability of ALERT
            p_danger: [B] float — sigmoid probability of DANGER
        """
        p_alert = torch.sigmoid(alert_logit)
        p_danger = torch.sigmoid(danger_logit)
        B = p_alert.shape[0]
        preds = torch.zeros(B, dtype=torch.long, device=p_alert.device)
        preds[p_danger > tau_danger] = 1   # OBSERVE
        preds[p_alert > tau_alert] = 2     # ALERT overrides OBSERVE
        return preds, p_alert, p_danger


class TrajectoryAwarePolicyHead(nn.Module):
    """
    Trajectory-Aware Policy Head — explicit per-timestep danger estimation
    with trajectory shape features for robust false alarm suppression.

    Key insight (Pjetri et al., ECCV-W 2024 extension):
      True collisions have monotonically increasing danger trajectories;
      false alarms / near-misses have NON-monotonic danger (rise then fall).
      OBSERVE acts as a sequential hypothesis test / confirmation buffer.
      Asymmetric monotonic constraint: enforce d(t)↑ only for ALERT; allow
      non-monotonic trajectories for OBSERVE.

    Architecture:
      Step 1: Per-timestep danger estimation
        belief[t] → proj(256) ⊕ tta_mean[t] ⊕ tta_var[t] → MLP(258→128→1) → σ → d[t]

      Step 2: Trajectory feature extraction (all differentiable)
        d_last, d_mean, d_max, d_gradient, d_acceleration, d_volatility, d_rise_ratio

      Step 3 (optional): GRU residual path for implicit temporal patterns

      Step 4: Classification
        [7 traj features ⊕ tta_last ⊕ tta_var_last (⊕ GRU_hidden)] → MLP → 3-class logits
    """

    def __init__(self, hidden_dim=2048, gru_hidden=256, n_actions=3,
                 dropout=0.2, use_gru=True):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.use_gru = use_gru
        self.n_actions = n_actions
        self.gru_hidden = gru_hidden

        # Step 1: per-timestep danger estimator
        self.belief_proj = nn.Linear(hidden_dim, 256)
        self.danger_estimator = nn.Sequential(
            nn.Linear(258, 128),          # 256 proj + 2 (tta_mean, tta_var)
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(128, 1),
        )
        # init danger output near 0 (sigmoid(0)=0.5) → slight negative bias
        nn.init.zeros_(self.danger_estimator[-1].weight)
        nn.init.constant_(self.danger_estimator[-1].bias, -0.5)

        # Step 3 (optional): GRU residual
        if use_gru:
            self.gru = nn.GRU(258, gru_hidden, num_layers=1,
                              batch_first=True, dropout=0)

        # Step 4: classifier
        # 7 trajectory features + 2 (tta_last, tta_var_last)
        clf_input_dim = 7 + 2
        if use_gru:
            clf_input_dim += gru_hidden

        self.classifier = nn.Sequential(
            nn.Linear(clf_input_dim, 128),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(128, 64),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(64, n_actions),
        )

    def forward(self, belief_seq, tta_mean_seq, tta_var_seq):
        """
        Args:
            belief_seq:   [B, T, hidden_dim]
            tta_mean_seq: [B, T]
            tta_var_seq:  [B, T]
        Returns:
            logits:   [B, n_actions]
            danger_t: [B, T] — per-timestep danger scores (for auxiliary loss)
        """
        B, T, _ = belief_seq.shape

        # Step 1: per-timestep danger
        proj = self.belief_proj(belief_seq)                        # [B, T, 256]
        tta_feat = torch.stack([tta_mean_seq, tta_var_seq], dim=-1)  # [B, T, 2]
        x = torch.cat([proj, tta_feat], dim=-1)                    # [B, T, 258]

        danger_t = torch.sigmoid(
            self.danger_estimator(x).squeeze(-1)                   # [B, T]
        )

        # Step 2: trajectory features (all differentiable)
        d_last = danger_t[:, -1]                                   # [B]
        d_mean = danger_t.mean(dim=1)                              # [B]
        d_max = danger_t.max(dim=1).values                         # [B]

        delta_d = danger_t[:, 1:] - danger_t[:, :-1]              # [B, T-1]
        d_gradient = delta_d.mean(dim=1)                           # [B]
        d_rise_ratio = (delta_d > 0).float().mean(dim=1)           # [B]

        if T > 2:
            d_volatility = delta_d.std(dim=1)                      # [B]
            delta2 = delta_d[:, 1:] - delta_d[:, :-1]             # [B, T-2]
            d_acceleration = delta2.mean(dim=1)                    # [B]
        else:
            d_volatility = torch.zeros(B, device=belief_seq.device)
            d_acceleration = torch.zeros(B, device=belief_seq.device)

        traj_features = torch.stack([
            d_last, d_mean, d_max, d_gradient,
            d_acceleration, d_volatility, d_rise_ratio,
        ], dim=-1)                                                 # [B, 7]

        # TTA context from last timestep
        tta_last = tta_mean_seq[:, -1].unsqueeze(-1)               # [B, 1]
        tta_var_last = tta_var_seq[:, -1].unsqueeze(-1)            # [B, 1]

        clf_input = torch.cat([traj_features, tta_last, tta_var_last], dim=-1)  # [B, 9]

        # Step 3 (optional): GRU residual
        if self.use_gru:
            _, h_n = self.gru(x)                                   # [1, B, gru_hidden]
            clf_input = torch.cat([clf_input, h_n.squeeze(0)], dim=-1)

        # Step 4: classification
        logits = self.classifier(clf_input)                        # [B, n_actions]
        return logits, danger_t


class TrajectoryAwarePOMDPHead(nn.Module):
    """Action-conditioned POMDP variant of TrajectoryAwarePolicyHead.

    Per-timestep belief update with explicit POMDP-style state transitions:
        h_t = GRU([belief_t ⊕ act_emb(prev_action_t) ⊕ tta_emb(tta_t)], h_{t-1})

    Outputs at each timestep:
        - logits_t [3]   per-step 3-class state (SILENT/OBSERVE/ALERT)
        - danger_t       per-step P(danger), kept for v7 monotonic-aux loss
        - tta_pred_t     per-step log-TTA reconstruction (auxiliary regularizer)

    Designed to be trained with teacher-forcing (`prev_action_t` =
    `action_label_seq[t-1]`); at inference time, can run autoregressively
    (use prev step's argmax as next prev_action) or with prev_action=SILENT
    init.
    """

    def __init__(self, hidden_dim=2560, gru_hidden=256, n_actions=3,
                 dropout=0.2, action_emb_dim=32, tta_emb_dim=32):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.gru_hidden = gru_hidden
        self.n_actions = n_actions

        # Action embedding (SILENT=0 / OBSERVE=1 / ALERT=2 + START=3 sentinel)
        self.action_emb = nn.Embedding(n_actions + 1, action_emb_dim)
        self.START_TOKEN = n_actions  # 3 = teacher-forcing start sentinel

        # TTA encoder (mean + var → embedding) — match shape of action_emb
        self.tta_encoder = nn.Sequential(
            nn.Linear(2, tta_emb_dim),
            nn.GELU(),
            nn.Linear(tta_emb_dim, tta_emb_dim),
        )

        # Belief projection to bring 2560 → 256
        self.belief_proj = nn.Linear(hidden_dim, 256)

        gru_input_dim = 256 + action_emb_dim + tta_emb_dim
        self.gru = nn.GRU(gru_input_dim, gru_hidden, num_layers=1,
                          batch_first=True, dropout=0)

        # Per-step state head (3-class)
        self.state_head = nn.Sequential(
            nn.Linear(gru_hidden, 128),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(128, n_actions),
        )

        # Per-step danger head (binary, for v7-style aux loss)
        self.danger_head = nn.Sequential(
            nn.Linear(gru_hidden, 64),
            nn.GELU(),
            nn.Linear(64, 1),
        )
        nn.init.zeros_(self.danger_head[-1].weight)
        nn.init.constant_(self.danger_head[-1].bias, -0.5)

        # Per-step TTA prediction head (log-TTA regression, for aux loss)
        self.tta_pred_head = nn.Sequential(
            nn.Linear(gru_hidden, 64),
            nn.GELU(),
            nn.Linear(64, 1),
        )

    def forward(self, belief_seq, tta_mean_seq, tta_var_seq,
                prev_action_seq=None):
        """
        Args:
            belief_seq:      [B, T, hidden_dim]
            tta_mean_seq:    [B, T]
            tta_var_seq:     [B, T]
            prev_action_seq: [B, T] long, prev_action_seq[t] = action at t-1
                             (teacher-forcing). If None, use START token.
        Returns:
            logits_seq:   [B, T, n_actions]    per-step state
            danger_seq:   [B, T]                per-step P(danger)
            tta_pred_seq: [B, T]                per-step log-TTA prediction
        """
        B, T, _ = belief_seq.shape

        if prev_action_seq is None:
            prev_action_seq = torch.full(
                (B, T), self.START_TOKEN, dtype=torch.long,
                device=belief_seq.device,
            )

        proj = self.belief_proj(belief_seq)                        # [B, T, 256]
        a_emb = self.action_emb(prev_action_seq)                   # [B, T, ae_dim]
        tta_feat = torch.stack(
            [tta_mean_seq, tta_var_seq], dim=-1)                    # [B, T, 2]
        tta_emb = self.tta_encoder(tta_feat)                       # [B, T, te_dim]

        x = torch.cat([proj, a_emb, tta_emb], dim=-1)              # [B, T, in]
        h_seq, _ = self.gru(x)                                     # [B, T, H]

        logits_seq = self.state_head(h_seq)                        # [B, T, 3]
        danger_seq = torch.sigmoid(
            self.danger_head(h_seq).squeeze(-1))                    # [B, T]
        tta_pred_seq = self.tta_pred_head(h_seq).squeeze(-1)       # [B, T] log-TTA
        return logits_seq, danger_seq, tta_pred_seq


class TemporalPolicyHead(nn.Module):
    """
    Temporal Belief Aggregation — GRU over K consecutive observation windows
    to capture danger escalation dynamics that single-frame beliefs miss.

    Motivation:
      - Single-frame AP locked at 0.24: beliefs separate dangerous/safe (AP=0.89)
        but cannot distinguish OBSERVE from ALERT.
      - Temporal gradient (danger increasing → ALERT vs stable → OBSERVE) requires
        multi-window context.

    Architecture:
      belief_seq [B, T, H] → Linear(H, 256) → concat(tta_mean, tta_var)
      → GRU(258, 256) → last hidden → MLP(256→128→3) → logits [B, 3]
    """

    def __init__(self, hidden_dim=2048, gru_hidden=256, n_actions=3, dropout=0.2):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.gru_hidden = gru_hidden

        self.belief_proj = nn.Linear(hidden_dim, 256)
        gru_input_dim = 256 + 2   # projected belief + tta_mean + tta_var
        self.gru = nn.GRU(gru_input_dim, gru_hidden, num_layers=1,
                          batch_first=True, dropout=0)
        self.head = nn.Sequential(
            nn.Linear(gru_hidden, 256),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(256, 128),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(128, n_actions),
        )

    def forward(self, belief_seq, tta_mean_seq, tta_var_seq):
        """
        Args:
            belief_seq:   [B, T, hidden_dim]
            tta_mean_seq: [B, T]
            tta_var_seq:  [B, T]
        Returns:
            logits: [B, n_actions]
        """
        proj = self.belief_proj(belief_seq)                          # [B, T, 256]
        tta = torch.stack([tta_mean_seq, tta_var_seq], dim=-1)       # [B, T, 2]
        x = torch.cat([proj, tta], dim=-1)                           # [B, T, 258]
        _, h_n = self.gru(x)                                         # [1, B, gru_hidden]
        return self.head(h_n.squeeze(0))                             # [B, 3]


# ═══════════════════════════════════════════════════════════════════════════════
# M10: Multi-Query PMA Aggregator (Pooling by Multi-head Attention)
#      Lee et al., "Set Transformer", ICML 2019 — universal set function approximator
# ═══════════════════════════════════════════════════════════════════════════════

class MultiQueryPMAAggregator(nn.Module):
    """
    K learnable query tokens cross-attend to per-frame belief tokens → K aggregated
    belief vectors that can specialise on orthogonal semantic axes (entity / motion
    / temporal / risk). Replaces mean_pool which collapses all frames to 1 vector.

    Input:
        beliefs_frame: [B, F, D]      per-frame beliefs (from per_frame cache)
        valid_mask:    [B, F] bool    True = valid frame, False = padded/missing
    Output:
        queries: [B, K, d_out]         K aggregated vectors
        attn:    [B, K, F]             attention weights (for interpretability/aux)
    """

    def __init__(
        self,
        d_in: int = 2048,
        d_out: int = 512,
        K: int = 4,
        n_heads: int = 4,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.K = K
        self.d_out = d_out

        # Learnable queries — one per semantic axis
        self.queries = nn.Parameter(torch.randn(1, K, d_out) * 0.02)

        self.in_proj = nn.Linear(d_in, d_out)
        self.mha = nn.MultiheadAttention(
            d_out, n_heads, dropout=dropout, batch_first=True,
        )
        self.ln1 = nn.LayerNorm(d_out)
        self.ffn = nn.Sequential(
            nn.Linear(d_out, d_out * 2),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_out * 2, d_out),
        )
        self.ln2 = nn.LayerNorm(d_out)

    def forward(self, beliefs_frame: torch.Tensor,
                valid_mask: torch.Tensor = None):
        B = beliefs_frame.shape[0]
        kv = self.in_proj(beliefs_frame.float())                  # [B, F, d_out]
        q = self.queries.expand(B, -1, -1).contiguous()           # [B, K, d_out]

        # key_padding_mask: True means *mask out* (invalid)
        kpm = None
        if valid_mask is not None:
            m = valid_mask.to(kv.device).bool()
            kpm = ~m
            # Guard against all-invalid rows (would give NaN in attention)
            all_invalid = kpm.all(dim=-1)
            if all_invalid.any():
                kpm = kpm.clone()
                kpm[all_invalid, 0] = False  # allow at least one slot

        attn_out, attn_w = self.mha(
            q, kv, kv,
            key_padding_mask=kpm,
            need_weights=True,
            average_attn_weights=True,
        )
        h = self.ln1(q + attn_out)
        h = self.ln2(h + self.ffn(h))
        return h, attn_w

    def orthogonality_loss(self) -> torch.Tensor:
        """L_ortho = ||Q Q^T - I||_F^2 / K^2   — prevents query collapse."""
        q = self.queries.squeeze(0)                                # [K, d_out]
        q = F.normalize(q, dim=-1)
        gram = q @ q.t()                                           # [K, K]
        eye = torch.eye(self.K, device=q.device, dtype=q.dtype)
        return ((gram - eye) ** 2).mean()


class MultiQueryPolicyHead(nn.Module):
    """
    Full M10 PolicyHead: aggregator + classifier.

    Pipeline:
      [B, F, D] per_frame beliefs
         → MultiQueryPMAAggregator → [B, K, d_out]
         → flatten [B, K*d_out]
         → concat (tta_mean, tta_var, prev_action embedding)
         → MLP → [B, 3]
    """

    def __init__(
        self,
        hidden_dim: int = 2048,
        d_out: int = 512,
        K: int = 4,
        n_heads: int = 4,
        n_actions: int = 3,
        dropout: float = 0.2,
    ):
        super().__init__()
        self.K = K
        self.d_out = d_out
        self.n_actions = n_actions

        self.aggregator = MultiQueryPMAAggregator(
            d_in=hidden_dim, d_out=d_out, K=K, n_heads=n_heads, dropout=0.1,
        )

        self.action_embedding = nn.Embedding(n_actions, 16)

        clf_input = K * d_out + 2 + 16
        self.classifier = nn.Sequential(
            nn.Linear(clf_input, 512),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(512, 256),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(256, n_actions),
        )

    def forward(
        self,
        beliefs_frame: torch.Tensor,    # [B, F, D]
        valid_mask: torch.Tensor,       # [B, F] bool
        tta_mean: torch.Tensor,         # [B]
        tta_var: torch.Tensor,          # [B]
        prev_action: torch.Tensor,      # [B] long
    ):
        agg, attn_w = self.aggregator(beliefs_frame, valid_mask)   # [B, K, d_out]
        flat = agg.reshape(agg.shape[0], -1)                        # [B, K*d_out]
        act_emb = self.action_embedding(prev_action)                # [B, 16]
        x = torch.cat([
            flat,
            tta_mean.unsqueeze(-1),
            tta_var.unsqueeze(-1),
            act_emb,
        ], dim=-1)
        logits = self.classifier(x)
        return logits, attn_w


class TransformerTemporalHead(nn.Module):
    """Transformer-based binary collision scorer over per-frame beliefs.

    Self-attention lets every frame pair interact directly, capturing patterns
    like "frame 7 looks dangerous vs frame 3 was safe" that sequential models
    (GRU) struggle with due to recency bias.

    Input:  beliefs_frame [B, T, 2048], tta_mean [B], tta_var [B]
    Output: binary logit [B]
    """

    def __init__(self, hidden_dim=2048, d_model=256, nhead=8, n_layers=2,
                 dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.frame_proj = nn.Sequential(
            nn.Linear(hidden_dim + 2, d_model),
            nn.LayerNorm(d_model),
        )
        self.cls_token = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)
        self.register_buffer('pe', self._sinusoidal_pe(65, d_model))
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead, dim_feedforward=d_model * 4,
            dropout=dropout, batch_first=True, activation='gelu',
        )
        self.encoder = nn.TransformerEncoder(encoder_layer,
                                             num_layers=n_layers)
        self.head = nn.Sequential(
            nn.Linear(d_model, 128),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(128, 1),
        )

    @staticmethod
    def _sinusoidal_pe(max_len, d_model):
        pe = torch.zeros(max_len, d_model)
        pos = torch.arange(max_len).unsqueeze(1).float()
        div = torch.exp(torch.arange(0, d_model, 2).float()
                        * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(pos * div)
        pe[:, 1::2] = torch.cos(pos * div)
        return pe.unsqueeze(0)

    def forward(self, beliefs_frame, tta_mean, tta_var, valid_mask=None):
        B, T, _ = beliefs_frame.shape
        tm = tta_mean.unsqueeze(1).unsqueeze(2).expand(B, T, 1)
        tv = tta_var.unsqueeze(1).unsqueeze(2).expand(B, T, 1)
        h = self.frame_proj(torch.cat([beliefs_frame, tm, tv], dim=-1))
        cls = self.cls_token.expand(B, -1, -1)
        h = torch.cat([cls, h], dim=1) + self.pe[:, :T + 1, :].to(h.device)
        pad_mask = None
        if valid_mask is not None:
            cls_valid = torch.ones(B, 1, dtype=torch.bool, device=h.device)
            pad_mask = ~torch.cat([cls_valid, valid_mask], dim=1)
        h = self.encoder(h, src_key_padding_mask=pad_mask)
        return self.head(h[:, 0, :]).squeeze(-1)


# ═══════════════════════════════════════════════════════════════════════════════
# M9: Spatial Attention Aggregator (for spatial4x4 cache)
#     Learnable query over 16 spatial cells per frame → per-frame belief;
#     then mean-over-F (or stack for downstream temporal model).
# ═══════════════════════════════════════════════════════════════════════════════

class SpatialAttentionAggregator(nn.Module):
    """
    Input:
        beliefs_grid:  [B, F, 16, D]   spatial4x4 cache
        valid_frames:  [B, F] bool
    Output:
        per_frame:     [B, F, d_out]   spatially attended per-frame belief
        frame_mean:    [B, d_out]      valid-frame mean of per_frame
        spatial_attn:  [B, F, 16]      spatial attention weights
    """

    def __init__(
        self,
        d_in: int = 2048,
        d_out: int = 512,
        n_heads: int = 4,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.d_out = d_out
        self.in_proj = nn.Linear(d_in, d_out)
        self.spatial_query = nn.Parameter(torch.randn(1, 1, d_out) * 0.02)
        self.mha = nn.MultiheadAttention(
            d_out, n_heads, dropout=dropout, batch_first=True,
        )
        self.ln = nn.LayerNorm(d_out)

    def forward(self, beliefs_grid: torch.Tensor, valid_frames: torch.Tensor):
        B, F_, S, D = beliefs_grid.shape
        x = self.in_proj(beliefs_grid.float())                     # [B, F, 16, d_out]
        # Flatten batch and frame for per-frame spatial attention
        x_flat = x.reshape(B * F_, S, self.d_out)                  # [B*F, 16, d_out]
        q = self.spatial_query.expand(B * F_, -1, -1).contiguous()

        attn_out, attn_w = self.mha(
            q, x_flat, x_flat, need_weights=True, average_attn_weights=True,
        )
        per_frame = self.ln(attn_out).squeeze(1)                   # [B*F, d_out]
        per_frame = per_frame.reshape(B, F_, self.d_out)           # [B, F, d_out]
        spatial_attn = attn_w.reshape(B, F_, S)                    # [B, F, 16]

        # Valid-frame mean pool (M9 single-belief output)
        valid = valid_frames.to(per_frame.device).float().unsqueeze(-1)  # [B, F, 1]
        denom = valid.sum(dim=1).clamp(min=1e-6)
        frame_mean = (per_frame * valid).sum(dim=1) / denom        # [B, d_out]

        return per_frame, frame_mean, spatial_attn


class SpatialPolicyHead(nn.Module):
    """
    Full M9 PolicyHead: spatial attention + classifier (single-belief output).
    Uses spatial4x4 cache. For a temporal variant, feed per_frame into GRU/PMA.
    """

    def __init__(
        self,
        hidden_dim: int = 2048,
        d_out: int = 512,
        n_heads: int = 4,
        n_actions: int = 3,
        dropout: float = 0.2,
    ):
        super().__init__()
        self.n_actions = n_actions
        self.aggregator = SpatialAttentionAggregator(
            d_in=hidden_dim, d_out=d_out, n_heads=n_heads, dropout=0.1,
        )
        self.action_embedding = nn.Embedding(n_actions, 16)

        clf_input = d_out + 2 + 16
        self.classifier = nn.Sequential(
            nn.Linear(clf_input, 512),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(512, 256),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(256, n_actions),
        )

    def forward(
        self,
        beliefs_grid: torch.Tensor,     # [B, F, 16, D]
        valid_frames: torch.Tensor,     # [B, F]
        tta_mean: torch.Tensor,
        tta_var: torch.Tensor,
        prev_action: torch.Tensor,
    ):
        _, frame_mean, spatial_attn = self.aggregator(beliefs_grid, valid_frames)
        act_emb = self.action_embedding(prev_action)
        x = torch.cat([
            frame_mean,
            tta_mean.unsqueeze(-1),
            tta_var.unsqueeze(-1),
            act_emb,
        ], dim=-1)
        logits = self.classifier(x)
        return logits, spatial_attn


class PatchTemporalHead(nn.Module):
    """Binary collision head over V-JEPA2 patch features.

    Input: patches [B, T, P, D]  (T=16 frames, P=256 patches, D=1024)
      1. Linear(D, hidden) projection per patch
      2. Spatial self-attention within each frame (1 layer, pooled via learnable CLS)
      3. Temporal self-attention across frame-level CLS summaries (2 layers)
      4. Temporal CLS → MLP → binary logit
    """

    def __init__(
        self,
        in_dim: int = 1024,
        hidden_dim: int = 256,
        n_spatial_layers: int = 1,
        n_temporal_layers: int = 2,
        n_heads: int = 4,
        dropout: float = 0.1,
        max_frames: int = 32,
    ):
        super().__init__()
        self.hidden_dim = hidden_dim

        self.proj = nn.Linear(in_dim, hidden_dim)

        self.spatial_cls = nn.Parameter(torch.zeros(1, 1, hidden_dim))
        nn.init.trunc_normal_(self.spatial_cls, std=0.02)

        spatial_layer = nn.TransformerEncoderLayer(
            d_model=hidden_dim,
            nhead=n_heads,
            dim_feedforward=hidden_dim * 4,
            dropout=dropout,
            batch_first=True,
            activation="gelu",
            norm_first=True,
        )
        self.spatial_encoder = nn.TransformerEncoder(spatial_layer, num_layers=n_spatial_layers)

        self.temporal_cls = nn.Parameter(torch.zeros(1, 1, hidden_dim))
        nn.init.trunc_normal_(self.temporal_cls, std=0.02)

        self.temporal_pos = nn.Parameter(torch.zeros(1, max_frames + 1, hidden_dim))
        nn.init.trunc_normal_(self.temporal_pos, std=0.02)

        temporal_layer = nn.TransformerEncoderLayer(
            d_model=hidden_dim,
            nhead=n_heads,
            dim_feedforward=hidden_dim * 4,
            dropout=dropout,
            batch_first=True,
            activation="gelu",
            norm_first=True,
        )
        self.temporal_encoder = nn.TransformerEncoder(temporal_layer, num_layers=n_temporal_layers)

        self.norm = nn.LayerNorm(hidden_dim)
        self.classifier = nn.Sequential(
            nn.Linear(hidden_dim, 128),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(128, 1),
        )

    def forward(self, patches: torch.Tensor) -> torch.Tensor:
        """patches: [B, T, P, D] → logits: [B]"""
        B, T, P, D = patches.shape
        assert T + 1 <= self.temporal_pos.shape[1], (
            f"T={T} exceeds max_frames; increase max_frames in PatchTemporalHead"
        )

        x = self.proj(patches)                              # [B, T, P, H]
        x = x.view(B * T, P, self.hidden_dim)

        cls = self.spatial_cls.expand(B * T, -1, -1)        # [B*T, 1, H]
        x = torch.cat([cls, x], dim=1)                      # [B*T, 1+P, H]
        x = self.spatial_encoder(x)
        frame_tokens = x[:, 0]                              # [B*T, H]
        frame_tokens = frame_tokens.view(B, T, self.hidden_dim)

        tcls = self.temporal_cls.expand(B, -1, -1)          # [B, 1, H]
        seq = torch.cat([tcls, frame_tokens], dim=1)        # [B, 1+T, H]
        seq = seq + self.temporal_pos[:, : 1 + T]
        seq = self.temporal_encoder(seq)

        clip = self.norm(seq[:, 0])                         # [B, H]
        return self.classifier(clip).squeeze(-1)            # [B]