"""
Sport-dataset dataloader for LED (soccer / football).

Schema mirrors NBADataset (data/dataloader_nba.py) so the trainer can reuse
its own collate and loop logic with minimal changes, but parameterized by:
  - data_dir       : absolute directory with train.npy / val.npy [/ test.npy]
  - num_agents     : agents per scene (23 for soccer/football)
  - split          : 'train' or 'val'   (val is used as the test set too)

The raw .npy files are shape (N, 30, A, 2). We split 10 past / 20 future
and return per-sample tensors in the same tuple layout as NBADataset so
the existing seq_collate logic in the sport trainer can produce the same
'pre_motion_3D'/'fut_motion_3D'/... dict.
"""

import os
import numpy as np
import torch
from torch.utils.data import Dataset


def sport_seq_collate(data):
    (pre_motion_3D, fut_motion_3D, pre_motion_mask, fut_motion_mask) = zip(*data)
    pre_motion_3D   = torch.stack(pre_motion_3D,   dim=0)
    fut_motion_3D   = torch.stack(fut_motion_3D,   dim=0)
    pre_motion_mask = torch.stack(pre_motion_mask, dim=0)
    fut_motion_mask = torch.stack(fut_motion_mask, dim=0)
    return {
        'pre_motion_3D':   pre_motion_3D,
        'fut_motion_3D':   fut_motion_3D,
        'fut_motion_mask': fut_motion_mask,
        'pre_motion_mask': pre_motion_mask,
        'traj_scale':      1,
        'pred_mask':       None,
        'seq':             'sport',
    }


class SportDataset(Dataset):
    """LED-style dataset for soccer/football raw .npy files."""

    def __init__(
        self,
        data_dir: str,
        num_agents: int,
        obs_len: int = 10,
        pred_len: int = 20,
        split: str = 'train',
    ):
        super().__init__()
        assert split in ('train', 'val'), f'unknown split {split}'
        self.obs_len   = obs_len
        self.pred_len  = pred_len
        self.seq_len   = obs_len + pred_len
        self.num_agents = num_agents

        fname = f'{split}.npy'
        path  = os.path.join(data_dir, fname)
        if not os.path.isfile(path):
            raise FileNotFoundError(f'missing {path}')

        trajs = np.load(path).astype(np.float32)   # (N, 30, A, 2)
        assert trajs.ndim == 4 and trajs.shape[1] == self.seq_len, \
            f'expected (N, {self.seq_len}, A, 2), got {trajs.shape}'
        assert trajs.shape[2] == num_agents, \
            f'expected {num_agents} agents, got {trajs.shape[2]}'

        self.data_len = len(trajs)
        print(f'[SportDataset] {split} {path} → {trajs.shape}')

        # [N, A, T, 2]  (LED convention: agent dim before time)
        self.traj_abs = torch.from_numpy(trajs).permute(0, 2, 1, 3).contiguous()

    def __len__(self):
        return self.data_len

    def __getitem__(self, index):
        pre  = self.traj_abs[index, :, :self.obs_len,  :]
        fut  = self.traj_abs[index, :, self.obs_len:, :]
        pre_mask = torch.ones(self.num_agents, self.obs_len)
        fut_mask = torch.ones(self.num_agents, self.pred_len)
        return [pre, fut, pre_mask, fut_mask]