"""MotionTransformerGraphV8 — V6 backbone with kinematic edge channels."""
from models.backbone_graph import MotionTransformerGraph
from models.graph_interaction_nba_v8 import FutureInteractionGraphV8


class MotionTransformerGraphV8(MotionTransformerGraph):

    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)

        self.future_graph = FutureInteractionGraphV8(
            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)

        params = sum(p.numel() for p in self.future_graph.parameters())
        logger.info(f"FutureInteractionGraphV8 (kinematic) parameters: {params:,}")