src/mira/data/cs2_stream.py

"""Streaming CS2-10k dataset for MIRA (no local dataset build).

RekaAI/CS2-10k is a WebDataset of first-person CS2 rounds. `index.parquet` lists every clip with:
  match_id, video_path, parquet_path, map, player_index, round_number, team, fps(48), total_time,
  width, height, match_num_clips, match_num_rounds, match_num_players, shard
Grouping key for MIRA's multi-perspective samples is (match_id, round_number): fixing both and
varying player_index gives up to `match_num_players` (=10) synchronized first-person views of the
same round — the CS2 analogue of MIRA's 4 Rocket League views.

This module streams those groups straight from the Hub and yields MIRA's own
`(VideoActionBatch, list[ClipMeta])` batches, so it drops into MIRA's trainer in place of
`create_loader` with zero on-disk dataset. It reuses MIRA's `decode_frames`, `KeyVocab`,
`ActionTensors`, `ClipMeta`, `VideoActionBatch`, and `_collate`.

Two important, honest caveats:
  * ALIGNMENT: CS2 gives no shared per-frame tick. Each round's per-player clips are 48fps and
    (empirically) start at round start, so we align at frame 0 and truncate a group to the shortest
    clip. Refine later with position/rotation cross-checks if needed.
  * SCALE / STREAMING UNIT: the 63TB main split is tar-sharded and a round's POVs are scattered
    across many shards, so complete groups are NOT co-located in one tar. This loader therefore
    fetches clips as INDIVIDUAL files by path (works today for the untarred `sample/` split — 3 full
    matches — and for any individual-file mirror). Training on the full tarred split at scale needs
    either re-sharding by (match,round) or a per-group multi-shard fetch; that's a separate step.
"""

from __future__ import annotations

import io
import random
from collections import defaultdict
from itertools import count
from typing import Any, Iterator

import numpy as np
import torch
from torch.utils.data import DataLoader, IterableDataset, get_worker_info

from huggingface_hub import hf_hub_download
import pandas as pd

import av  # ffmpeg-backed decode (works in this venv; avoids the torchcodec dependency)

from .actions import KeyVocab
from .batch import VideoActionBatch
from .clips import compute_stride
from .training_loader import ClipMeta, _collate  # reuse MIRA's collate + meta type


def _decode_av(mp4_bytes: bytes, frame_indices: list[int], frame_size: tuple[int, int] | None):
    """Decode the given source-frame indices from mp4 bytes to (T, C, H, W) uint8, via PyAV."""
    want = set(frame_indices)
    hi = max(frame_indices)
    grabbed: dict[int, np.ndarray] = {}
    with av.open(io.BytesIO(mp4_bytes)) as container:
        for i, frame in enumerate(container.decode(container.streams.video[0])):
            if i in want:
                grabbed[i] = frame.to_ndarray(format="rgb24")  # (H, W, 3) uint8
            if i >= hi:
                break
    arr = np.stack([grabbed[i] for i in frame_indices])          # (T, H, W, 3)
    t = torch.from_numpy(arr).permute(0, 3, 1, 2).contiguous()   # (T, C, H, W) uint8
    if frame_size is not None:
        t = torch.nn.functional.interpolate(
            t.float(), size=frame_size, mode="bilinear", align_corners=False
        ).clamp(0, 255).to(torch.uint8)
    return t

# CS2 held-key characters (per-frame `actions` string), stable multi-hot order. '-' == no input.
CS2_KEYS: tuple[str, ...] = ("W", "A", "S", "D", "J", "C", "R", "V", "[", "]")

REPO_ID = "RekaAI/CS2-10k"


def _rank_world() -> tuple[int, int]:
    if torch.distributed.is_available() and torch.distributed.is_initialized():
        return torch.distributed.get_rank(), torch.distributed.get_world_size()
    return 0, 1


def _fetch_bytes(path: str) -> bytes:
    """Fetch an INDIVIDUAL file by repo path (works for the untarred `sample/` split)."""
    local = hf_hub_download(REPO_ID, path, repo_type="dataset")
    with open(local, "rb") as f:
        return f.read()


# --- ranged tar-member fetch for the FULL (tarred) split ------------------------------------------
# A round's POVs are scattered across many ~2GB tars. We open each tar over HfFileSystem (a seekable,
# HTTP-range-backed file) with tarfile mode="r", which walks the member headers via small ranged
# reads ONCE per shard, then extracts a single member with one ranged read of just its bytes — so we
# never download a whole 2GB tar. Open TarFiles are cached per shard.
import tarfile

_HF_FS = None
_TAR_CACHE: dict[str, tuple] = {}


def _hf_fs():
    global _HF_FS
    if _HF_FS is None:
        from huggingface_hub import HfFileSystem
        _HF_FS = HfFileSystem()
    return _HF_FS


def _open_shard(shard: str):
    if shard not in _TAR_CACHE:
        f = _hf_fs().open(f"datasets/{REPO_ID}/{shard}", "rb")   # seekable, range-backed
        tf = tarfile.open(fileobj=f, mode="r")                    # header walk (ranged) once
        members = {m.name.rsplit("/", 1)[-1]: m for m in tf.getmembers()}
        if len(_TAR_CACHE) > 24:                                  # bound open handles
            old = next(iter(_TAR_CACHE))
            try:
                _TAR_CACHE.pop(old)[0].close()
            except Exception:
                pass
        _TAR_CACHE[shard] = (tf, members)
    return _TAR_CACHE[shard]


def _fetch_from_tar(shard: str, repo_path: str) -> bytes:
    tf, members = _open_shard(shard)
    m = members[repo_path.rsplit("/", 1)[-1]]
    return tf.extractfile(m).read()


def _build_action_arrays(
    frame_data: np.ndarray, vocab: KeyVocab, stride: int
) -> tuple[torch.Tensor, torch.Tensor]:
    """From a clip's 48fps per-frame records build downsampled (n_steps, n_keys) int32 multi-hot key
    presses (OR-ed over each stride window) and (n_steps, 2) float32 mean mouse deltas."""
    n_keys = len(vocab)
    n_steps = len(frame_data) // stride
    keys = torch.zeros((n_steps, n_keys), dtype=torch.int32)
    mouse = torch.zeros((n_steps, 2), dtype=torch.float32)
    for s in range(n_steps):
        mx = my = 0.0
        for j in range(s * stride, (s + 1) * stride):
            rec = frame_data[j]
            a = rec["actions"]
            if a and a != "-":
                for ch in a:
                    idx = vocab._index.get(ch)
                    if idx is not None:
                        keys[s, idx] = 1
            mx += float(rec["mouse_x_delta"])
            my += float(rec["mouse_y_delta"])
        mouse[s, 0] = mx / stride
        mouse[s, 1] = my / stride
    return keys, mouse


class Cs2StreamingDataset(IterableDataset):
    """Streams (match_id, round_number) POV groups from CS2-10k as MIRA per-perspective samples.

    Yields dicts shaped exactly like MIRA's training_loader `_decode_sample` output
    ({"video","actions","metadata"}), so `_collate` turns a batch into `(VideoActionBatch, [ClipMeta])`.
    """

    def __init__(
        self,
        index: pd.DataFrame,
        *,
        vocab: KeyVocab,
        clip_len: int = 16,
        target_fps: int = 16,
        source_fps: int = 48,
        n_players: int = 4,
        frame_size: tuple[int, int] | None = None,
        mouse_sensitivity: float = 1.0,
        shuffle: bool = True,
        infinite: bool = True,
        seed: int = 2025,
        from_tar: bool = False,
    ) -> None:
        super().__init__()
        self.index = index
        self.from_tar = from_tar
        self.vocab = vocab
        self.clip_len = clip_len
        self.target_fps = target_fps
        self.source_fps = source_fps
        self.n_players = n_players
        self.frame_size = frame_size
        self.mouse_sensitivity = mouse_sensitivity
        self.shuffle = shuffle
        self.infinite = infinite
        self.seed = seed
        self.stride = compute_stride(source_fps, target_fps)
        # group rows into (match_id, round_number) -> [row, ...] sorted by player_index
        groups: dict[tuple[str, int], list[dict]] = defaultdict(list)
        for r in index.to_dict("records"):
            groups[(r["match_id"], int(r["round_number"]))].append(r)
        # keep only groups with enough perspectives to form an n_players block
        self.groups: list[list[dict]] = [
            sorted(v, key=lambda r: r["player_index"])
            for v in groups.values()
            if len(v) >= n_players
        ]

    def _my_groups(self) -> list[list[dict]]:
        rank, world = _rank_world()
        info = get_worker_info()
        wid, nw = (info.id, info.num_workers) if info else (0, 1)
        return self.groups[rank::world][wid::nw]

    def _load_group(self, rows: list[dict]) -> tuple[list[bytes], list[np.ndarray], list[dict]]:
        import io as _io
        vids, fdata, meta = [], [], []
        for r in rows[: self.n_players]:
            if self.from_tar:                                        # full split: ranged tar members
                vids.append(_fetch_from_tar(r["shard"], r["video_path"]))
                pq = pd.read_parquet(_io.BytesIO(_fetch_from_tar(r["shard"], r["parquet_path"])))
            else:                                                    # sample split: individual files
                vids.append(_fetch_bytes(r["video_path"]))
                pq = pd.read_parquet(hf_hub_download(REPO_ID, r["parquet_path"], repo_type="dataset"))
            fdata.append(pq.iloc[0]["frame_data"])
            meta.append(r)
        return vids, fdata, meta

    def _samples_from_group(self, rows: list[dict]) -> Iterator[dict[str, Any]]:
        vids, fdata, meta = self._load_group(rows)
        n_src = min(len(fd) for fd in fdata)              # align at frame 0, truncate to shortest
        n_steps = n_src // self.stride
        if n_steps < self.clip_len:
            return
        # per-perspective full-round action arrays, then slice per clip window
        key_arrs, mouse_arrs = [], []
        for fd in fdata:
            k, m = _build_action_arrays(fd[: n_steps * self.stride], self.vocab, self.stride)
            key_arrs.append(k)
            mouse_arrs.append(m)
        # Decode every target-fps frame of the clip ONCE per perspective, then slice windows from it.
        # (Decoding per window re-decoded from the start of the mp4 each time -> O(n^2) and starved
        #  the GPU; this is a single linear pass.)
        all_idx = [s * self.stride for s in range(n_steps)]
        full = [_decode_av(vids[p], all_idx, self.frame_size) for p in range(self.n_players)]
        clip_id = 0
        for start in range(0, n_steps - self.clip_len + 1, self.clip_len):
            frame_indices = all_idx[start : start + self.clip_len]
            for p in range(self.n_players):
                video = full[p][start : start + self.clip_len]  # (T,C,H,W) uint8, already decoded
                act = _make_action_tensors(
                    key_arrs[p][start : start + self.clip_len],
                    mouse_arrs[p][start : start + self.clip_len],
                    self.vocab,
                    self.mouse_sensitivity,
                )
                yield {
                    "video": video,
                    "actions": act,
                    "metadata": ClipMeta(
                        match_id=meta[p]["match_id"],
                        perspective=p,
                        player_id=int(meta[p]["player_index"]),
                        clip_id=clip_id,
                        chunk_idx=int(meta[p]["round_number"]),
                        frame_indices=list(frame_indices),
                    ),
                }
            clip_id += 1

    def __iter__(self) -> Iterator[dict[str, Any]]:
        rank, _ = _rank_world()
        info = get_worker_info()
        rng = random.Random(self.seed + rank * 1024 + (info.id if info else 0))
        for _epoch in count() if self.infinite else range(1):
            order = self._my_groups()
            if self.shuffle:
                rng.shuffle(order)
            for rows in order:
                try:
                    yield from self._samples_from_group(rows)
                except Exception as e:  # a bad clip shouldn't kill the epoch
                    print(f"[cs2_stream] skipping group {rows[0]['match_id']}: {e}")


def _make_action_tensors(keys: torch.Tensor, mouse: torch.Tensor, vocab: KeyVocab, sens: float):
    """Build a MIRA ActionTensors (batch=1) from a clip window's key/mouse tensors."""
    from mira.world_model.actions_config import ActionConfig, ActionTensors

    cfg = ActionConfig(valid_keys=list(vocab.keys))
    at = ActionTensors(config=cfg, batch_size=1)
    at.key_presses = keys.unsqueeze(0).to(torch.int32)          # (1, T, n_keys)
    at.mouse_movements = mouse.unsqueeze(0).to(torch.float32)    # (1, T, 2)
    at.game_mouse_sensitivity = torch.full((1,), float(sens), dtype=torch.float32)
    return at


def load_cs2_index(subset: str = "sample") -> pd.DataFrame:
    """Return the CS2 index as a DataFrame. `subset='sample'` uses the untarred 3-match `sample/`
    split whose clips are individually fetchable (recommended for streaming today); `subset='full'`
    uses the root index (tarred; see the module docstring caveat before using at scale)."""
    if subset == "sample":
        p = hf_hub_download(REPO_ID, "sample/index.parquet", repo_type="dataset")
        df = pd.read_parquet(p)
        # sample/ paths are relative to sample/; make them repo-relative
        for col in ("video_path", "parquet_path"):
            if col in df.columns:
                df[col] = df[col].apply(lambda x: x if x.startswith("sample/") else f"sample/{x}")
        return df
    p = hf_hub_download(REPO_ID, "index.parquet", repo_type="dataset")
    return pd.read_parquet(p)


def create_cs2_loader(
    *,
    subset: str = "sample",
    index: pd.DataFrame | None = None,
    maps: list[str] | None = None,
    clip_len: int = 16,
    target_fps: int = 16,
    n_players: int = 4,
    batch_size: int = 4,
    num_workers: int = 0,
    frame_size: tuple[int, int] | None = None,
    valid_keys: list[str] | None = None,
    shuffle: bool = True,
    infinite: bool = True,
    seed: int = 2025,
    prefetch_factor: int = 2,
) -> DataLoader:
    """Drop-in replacement for MIRA's `create_loader`, streaming CS2-10k from the Hub (no disk).

    Yields `(VideoActionBatch, list[ClipMeta])`. `n_players` perspectives per (match, round) are
    grouped contiguously so the collate stacks them into one multi-perspective block, exactly as
    MIRA's own loader does.
    """
    if index is None:
        index = load_cs2_index(subset)
    if maps:
        index = index[index["map"].isin(maps)]
    vocab = KeyVocab(tuple(valid_keys) if valid_keys else CS2_KEYS, on_unknown="ignore")
    ds = Cs2StreamingDataset(
        index,
        vocab=vocab,
        clip_len=clip_len,
        target_fps=target_fps,
        n_players=n_players,
        frame_size=frame_size,
        shuffle=shuffle,
        infinite=infinite,
        seed=seed,
        from_tar=(subset == "full"),   # full split = ranged tar-member fetch; sample = individual files
    )
    return DataLoader(
        ds,
        batch_size=batch_size * n_players,   # n_players contiguous rows per sample-block
        num_workers=num_workers,
        collate_fn=_collate,
        drop_last=True,
        prefetch_factor=prefetch_factor if num_workers else None,
    )