"""
MotionTransformerGraphV16 — Sandwich: graph BEFORE and AFTER A-attn.

K-attn → graph_pre → A-attn → graph_post → context fusion

Pre-graph primes A-attn with future interaction context.
Post-graph refines the A-attn output with updated interaction.
Two separate graph instances (different weights, same architecture).
"""

import torch
import torch.nn as nn
from einops import rearrange, repeat

from models.backbone_graph import MotionTransformerGraph
from models.graph_interaction_nba_v6 import FutureInteractionGraphV6


class MotionTransformerGraphV16(MotionTransformerGraph):
    """Graph sandwich around A-attn."""

    def __init__(self, model_config, logger, config,
                 graph_num_gnn_layers=2, graph_dropout=0.1,
                 top_n_neighbors=5, rel_traj_hidden=32, y0_score_dim=32):
        super().__init__(model_config, logger, config,
                         graph_num_gnn_layers=graph_num_gnn_layers,
                         graph_dropout=graph_dropout)

        # Pre-A-attn graph (primes attention)
        self.graph_pre = FutureInteractionGraphV6(
            embed_dim=self.dim, future_steps=self.T_future,
            num_agents=self.A, num_heads=4, dropout=graph_dropout,
            num_gnn_layers=1,  # lighter: 1 layer
            time_dim=self.dim,
            top_n_neighbors=top_n_neighbors,
            rel_traj_hidden=rel_traj_hidden, y0_score_dim=y0_score_dim)

        # Post-A-attn graph (refines) — this IS the parent's self.future_graph
        self.future_graph = FutureInteractionGraphV6(
            embed_dim=self.dim, future_steps=self.T_future,
            num_agents=self.A, num_heads=4, dropout=graph_dropout,
            num_gnn_layers=graph_num_gnn_layers,
            time_dim=self.dim,
            top_n_neighbors=top_n_neighbors,
            rel_traj_hidden=rel_traj_hidden, y0_score_dim=y0_score_dim)

        p1 = sum(p.numel() for p in self.graph_pre.parameters())
        p2 = sum(p.numel() for p in self.future_graph.parameters())
        logger.info(f"V16: pre-graph: {p1:,}, post-graph: {p2:,}")

    def _forward_impl(self, y, time, x_data,
                      y_0_for_graph=None, sigma_for_graph=None,
                      skip_graph=False):
        if y.size(-1) == 2:
            y = y.reshape((-1, self.model_cfg.NUM_PROPOSED_QUERY,
                           self.A, self.T_future * 2))
        device = y.device
        B, K, A, _ = y.shape

        encoder_out = self.context_encoder(x_data['past_traj_original_scale'])
        encoder_out_batch = repeat(encoder_out, 'b a d -> b k a d', k=K, a=A)

        y_emb = self.noisy_y_mlp(y)

        time_ = time
        if self.config.denoising_method == 'fm':
            time = time * 1000.0
        t_emb = self.time_mlp(time)
        t_emb_batch = repeat(t_emb, 'b d -> b k a d', b=B, k=K, a=A)

        k_pe = self.motion_query_embedding(
            torch.arange(self.model_cfg.NUM_PROPOSED_QUERY, device=device))
        k_pe_batch = repeat(k_pe, 'k d -> b k a d', b=B, a=A)
        a_pe = self.agent_order_embedding(
            torch.arange(self.model_cfg.CONTEXT_ENCODER.NUM_OF_ATTN_NEIGHBORS, device=device))
        a_pe_batch = repeat(a_pe, 'a d -> b k a d', b=B, k=K)

        # K-attn
        y_emb_k = rearrange(self.apply_PE(y_emb, k_pe_batch, a_pe_batch), 'b k a d -> (b a) k d')
        y_emb_k = self.noisy_y_attn_k(y_emb_k)
        y_emb = rearrange(y_emb_k, '(b a) k d -> b k a d', b=B, a=A)

        # Compute geometry once
        y_abs = None
        if not skip_graph:
            y_graph_src = (y_0_for_graph.view(B, K, A, self.T_future, 2)
                           if y_0_for_graph is not None
                           else y.view(B, K, A, self.T_future, 2))
            y_graph_unnorm = self._unnormalize_y(y_graph_src)
            init_pos = x_data['past_traj_original_scale'][:, :, -1, :2]
            y_abs = y_graph_unnorm + init_pos.unsqueeze(1).unsqueeze(3)
            tau = time_

            # ---- PRE-GRAPH (before A-attn) ----
            y_emb = self.graph_pre(y_emb, y_abs, t_emb, tau, sigma_agent=None)

        # ---- A-attn ----
        y_emb_a = rearrange(y_emb, 'b k a d -> (b k) a d')
        y_emb_a = self.noisy_y_attn_a(y_emb_a)
        y_emb = rearrange(y_emb_a, '(b k) a d -> b k a d', b=B, k=K)

        # Embedding dropout
        if self.training and self.config.get('drop_method', None) == 'emb':
            m, k_drop = self.config.drop_logi_m, self.config.drop_logi_k
            p_m = 1 / (1 + torch.exp(-k_drop * (time_ - m)))
            p_m = p_m[:, None, None, None]
            y_emb = y_emb.masked_fill(torch.rand_like(p_m) < p_m, 0.)

        # ---- POST-GRAPH (after A-attn, same as V6) ----
        if not skip_graph and y_abs is not None:
            y_emb = self.future_graph(
                y_emb, y_abs, t_emb, tau, sigma_agent=sigma_for_graph)

        # Context fusion + decoder
        emb_fusion = self.init_emb_fusion_mlp(
            torch.cat((encoder_out_batch, y_emb, t_emb_batch), dim=-1))
        query_token = self.post_pe_cat_mlp(
            self.apply_PE(emb_fusion, k_pe_batch, a_pe_batch))
        readout_token = self.motion_decoder(query_token, t_emb)

        denoiser_x = self.reg_head(readout_token)
        denoiser_cls = self.cls_head(readout_token).squeeze(-1)
        logvar = self.logvar_head(readout_token)

        return denoiser_x, denoiser_cls, logvar