"""
MotionTransformerGraphV5 — two-pass backbone with uncertainty-aware semantic
sparse graph (FutureInteractionGraphV5).
"""

from models.backbone_graph import MotionTransformerGraph
from models.graph_interaction_nba_v5 import FutureInteractionGraphV5


class MotionTransformerGraphV5(MotionTransformerGraph):
    """MotionTransformerGraph with uncertainty-aware semantic sparse graph.

    Extra constructor kwargs:
        top_n_neighbors (int, default 5)
        rel_traj_hidden (int, default 32)
        y0_score_dim    (int, default 32)
    """

    def __init__(self, model_config, logger, config,
                 graph_num_gnn_layers: int = 2,
                 graph_dropout: float = 0.1,
                 top_n_neighbors: int = 5,
                 rel_traj_hidden: int = 32,
                 y0_score_dim: 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

        self.future_graph = FutureInteractionGraphV5(
            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,
            top_n_neighbors = top_n_neighbors,
            rel_traj_hidden = rel_traj_hidden,
            y0_score_dim    = y0_score_dim,
        )

        params_graph = sum(p.numel() for p in self.future_graph.parameters())
        logger.info("FutureInteractionGraphV5 parameters: {:,}".format(params_graph))
        logger.info("Top-N neighbors: {:d} / {:d}".format(top_n_neighbors, self.A - 1))