"""MotionTransformerGraphV9b — V6 + adaptive readout (3ch, no kinematic)."""
from models.backbone_graph import MotionTransformerGraph
from models.graph_interaction_nba_v9b import FutureInteractionGraphV9b

class MotionTransformerGraphV9b(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 = FutureInteractionGraphV9b(
            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"FutureInteractionGraphV9b (adaptive readout) parameters: {params:,}")