MoFlow/data/dataloader_sdd_moflow.py

"""SDD dataloader for MoFlow with MULTI-SCENE BATCHING via per-batch padding.

Strategy:
  - Sort scenes by A (number of agents) at init time.
  - BatchSampler builds mini-batches from contiguous A-sorted indices so that
    all scenes in a batch have similar A → padding waste is small (≤ 2×).
  - Collate pads each batch to its own max_A and returns an `agent_mask`
    [B, max_A] (True = real agent, False = padding) that the backbone and
    graph module use to ignore padded slots.

Scenes within a batch attend ONLY to agents in their own scene via a
block-diagonal `traj_mask`; padded slots are zeroed out in the mask so
no information leaks across scenes or into padding.
"""
import os, pickle, math, random
import numpy as np
import torch
from torch.utils.data import Dataset, Sampler

from utils.normalization import normalize_min_max


def _sdd_to_scene(sdd, obs_len, pred_len):
    """Convert one raw scene [A, 20, 2] numpy array to torch tensors needed
    by MoFlow: (past_6ch, fut_rel, fut_vel, init_pos).
    Coords are already /50 and mean-centered per scene from preprocessing."""
    traj = torch.from_numpy(sdd).float()        # [A, 20, 2]
    pre  = traj[:, :obs_len].contiguous()
    fut  = traj[:, obs_len:].contiguous()
    init = pre[:, -1:]
    abs_ch = pre
    rel_ch = pre - init
    vel_ch = torch.cat([rel_ch[:, 1:] - rel_ch[:, :-1],
                        torch.zeros_like(rel_ch[:, -1:])], dim=1)
    past_6 = torch.cat([abs_ch, rel_ch, vel_ch], dim=-1)   # [A, T_past, 6]
    fut_r  = fut - init                                     # [A, T_future, 2]
    fut_vel = torch.cat([fut_r[:, 1:] - fut_r[:, :-1],
                         torch.zeros_like(fut_r[:, -1:])], dim=1)
    return past_6, fut_r, fut_vel


class SortedByASampler(Sampler):
    """Yield single indices in A-sorted order with bucket-local shuffling.
    Use with DataLoader(..., sampler=this, batch_size=B): consecutive indices
    land in the same mini-batch so batches end up A-similar. Bucket size W =
    4 * batch_size controls how much A-shuffling happens per epoch.
    """
    def __init__(self, a_counts, batch_size, shuffle=True, seed=0):
        self.a_counts = np.asarray(a_counts)
        self.batch_size = int(batch_size)
        self.shuffle = shuffle
        self.seed = seed
        self.epoch = 0

    def set_epoch(self, epoch):
        self.epoch = int(epoch)

    def __iter__(self):
        order = np.argsort(self.a_counts, kind='stable')
        if self.shuffle:
            rng = np.random.RandomState(self.seed + self.epoch)
            W = 4 * self.batch_size
            for start in range(0, len(order), W):
                end = min(start + W, len(order))
                rng.shuffle(order[start:end])
        for idx in order.tolist():
            yield int(idx)

    def __len__(self):
        return len(self.a_counts)


def collate_padded(batch):
    """Pad a list of `(past_6, fut_rel, fut_vel)` scenes (each [A_i, T, *]) to
    a common max_A in this batch, stack, and emit an agent_mask.

    Returns the same dict structure as dataloader_nba's seq_collate_nba plus
    extra keys: 'agent_mask' [B, max_A] and 'A_per_scene' [B].
    """
    B = len(batch)
    A_list = [x[0].shape[0] for x in batch]
    max_A = max(A_list)
    T_past = batch[0][0].shape[1]
    T_fut  = batch[0][1].shape[1]

    past_pad = torch.zeros(B, max_A, T_past, 6)
    fut_pad  = torch.zeros(B, max_A, T_fut,  2)
    vel_pad  = torch.zeros(B, max_A, T_fut,  2)
    agent_mask = torch.zeros(B, max_A, dtype=torch.bool)   # True = real

    for i, (p6, fr, fv) in enumerate(batch):
        A_i = p6.shape[0]
        past_pad[i, :A_i] = p6
        fut_pad[i, :A_i]  = fr
        vel_pad[i, :A_i]  = fv
        agent_mask[i, :A_i] = True

    # Block-diagonal traj_mask zeroed on padding
    N = B * max_A
    traj_mask = torch.zeros(N, N)
    flat_mask = agent_mask.view(-1)    # [B*max_A]
    for i in range(B):
        s, e = i * max_A, (i + 1) * max_A
        block = agent_mask[i].unsqueeze(0) & agent_mask[i].unsqueeze(1)  # [max_A, max_A]
        traj_mask[s:e, s:e] = block.float()

    return {
        'batch_size': torch.tensor(B),
        'past_traj':           past_pad,                  # [B, max_A, T_past, 6]
        'fut_traj':            fut_pad,                   # [B, max_A, T_fut,  2]
        'past_traj_original_scale': past_pad.clone(),     # same as past_traj (normalize afterwards if needed)
        'fut_traj_original_scale':  fut_pad.clone(),
        'traj_mask': traj_mask,                           # [B*max_A, B*max_A]
        'fut_traj_vel': vel_pad,                          # [B, max_A, T_fut, 2]
        'agent_mask': agent_mask,                         # [B, max_A] True=real
    }


class SDDDatasetMinMax(Dataset):
    def __init__(
        self,
        obs_len=8, pred_len=12,
        training=True,
        cfg=None,
        data_dir='',
        split='train',
        data_norm='min_max',
    ):
        super().__init__()
        self.obs_len = obs_len
        self.pred_len = pred_len
        self.cfg = cfg
        self.data_norm = data_norm
        self.training = training

        path = os.path.join(data_dir, f'sdd_{split}_v2.pkl')
        with open(path, 'rb') as f:
            self.raw_scenes = pickle.load(f)

        # Pre-compute A per scene for the BatchSampler
        self.a_counts = [int(s.shape[0]) for s in self.raw_scenes]
        print(f'[SDDDatasetMinMax] {split}: {len(self.raw_scenes)} scenes, '
              f'A min/mean/max = {min(self.a_counts)}/{np.mean(self.a_counts):.1f}/{max(self.a_counts)}')

        # Compute global min/max for min_max normalization (on training set)
        if training and data_norm == 'min_max':
            all_past, all_fut = [], []
            for s in self.raw_scenes:
                past_6, fut_r, _ = _sdd_to_scene(s, obs_len, pred_len)
                all_past.append(past_6); all_fut.append(fut_r)
            past_cat = torch.cat(all_past, dim=0)
            fut_cat  = torch.cat(all_fut,  dim=0)
            cfg.past_traj_max = past_cat.max()
            cfg.past_traj_min = past_cat.min()
            cfg.fut_traj_max  = fut_cat.max()
            cfg.fut_traj_min  = fut_cat.min()

    def __len__(self):
        return len(self.raw_scenes)

    def __getitem__(self, idx):
        past_6, fut_r, fut_vel = _sdd_to_scene(self.raw_scenes[idx], self.obs_len, self.pred_len)
        if self.data_norm == 'min_max':
            past_6 = normalize_min_max(past_6, self.cfg.past_traj_min, self.cfg.past_traj_max, -1, 1).contiguous()
            fut_r  = normalize_min_max(fut_r,  self.cfg.fut_traj_min,  self.cfg.fut_traj_max,  -1, 1).contiguous()
        return past_6, fut_r, fut_vel


# Alias kept for backwards-compat with import lines in fm_sdd.py / fm_sdd_graph.py
seq_collate_sdd = collate_padded