"""
MotionTransformerGraphV3 — two-pass backbone with spatio-temporal graph edges.

Identical to MotionTransformerGraphV2 except it uses FutureInteractionGraphV3,
which encodes the full relative trajectory sequence (not just mean position)
using a small transformer (RelTrajEncoder) with self-attention over T timesteps.
"""

from models.backbone_graph import MotionTransformerGraph
from models.graph_interaction_nba_v3 import FutureInteractionGraphV3


class MotionTransformerGraphV3(MotionTransformerGraph):
    """MotionTransformerGraph with spatio-temporal future interaction edges.

    Constructor arguments identical to MotionTransformerGraph, plus:
        rel_traj_hidden (int, default 32): hidden dim inside RelTrajEncoder.
    """

    def __init__(self, model_config, logger, config,
                 graph_num_gnn_layers: int = 2,
                 graph_dropout: float = 0.1,
                 rel_traj_hidden: int = 32):
        super().__init__(model_config, logger, config,
                         graph_num_gnn_layers=graph_num_gnn_layers,
                         graph_dropout=graph_dropout)

        D        = self.dim
        time_dim = D

        # Replace the V1/V2 graph module with V3 (spatio-temporal edge encoder)
        self.future_graph = FutureInteractionGraphV3(
            embed_dim       = D,
            future_steps    = self.T_future,
            num_agents      = self.A,
            num_heads       = 4,
            dropout         = graph_dropout,
            num_gnn_layers  = graph_num_gnn_layers,
            time_dim        = time_dim,
            rel_traj_hidden = rel_traj_hidden,
        )

        params_graph = sum(p.numel() for p in self.future_graph.parameters())
        logger.info("FutureInteractionGraphV3 parameters: {:,}".format(params_graph))