src/data/datasets.py

"""
MultiSense-DF — PyTorch Dataset Implementations
================================================
Provides:
  - MultiModalDataset : generic real/fake folder structure  (real/ + fake/ subfolders)
  - DFDCDataset       : DFDC Preview dataset with metadata.json labels

Both datasets return a dict:
  {
    'frames':      (T, 3, 224, 224)  float32,
    'waveform':    (sr*duration,)    float32,
    'mouth_crops': (T, 3,  96,  96)  float32,
    'mel_specs':   (T, 1,  80,  W)   float32,
    'label':       int  (0=real, 1=fake),
    'video_path':  str,
  }
"""

from __future__ import annotations

import json
import os
import random
import shutil
from pathlib import Path
from typing import Callable, Dict, List, Optional, Tuple

import torch
from torch.utils.data import Dataset

from .preprocessing import (
    extract_frames,
    extract_audio_waveform,
    extract_mel_spectrogram,
    extract_mouth_crops,
    VISUAL_TRAIN_TRANSFORM,
    VISUAL_VAL_TRANSFORM,
)


# ---------------------------------------------------------------------------
# Helper
# ---------------------------------------------------------------------------

_VIDEO_EXTS: frozenset = frozenset({".mp4", ".avi", ".mov", ".mkv", ".webm", ".flv"})


def _scan_videos(directory: Path) -> List[Path]:
    """Recursively collect all video files under *directory*."""
    return sorted(
        p for p in directory.rglob("*") if p.suffix.lower() in _VIDEO_EXTS
    )


# ---------------------------------------------------------------------------
# MultiModalDataset
# ---------------------------------------------------------------------------

class MultiModalDataset(Dataset):
    """
    Generic dataset for any folder that contains 'real' and 'fake' subfolders.

    Expected layout::

        root_dir/
          real/
            video001.mp4
            video002.mp4
            ...
          fake/
            video003.mp4
            video004.mp4
            ...

    Parameters
    ----------
    root_dir : str | Path
        Root directory containing 'real' and 'fake' sub-folders.
    split : str
        One of {'train', 'val', 'test'}.  Controls augmentation and which
        cached split is used when *cache_dir* is set.
    num_frames : int
        Number of frames to uniformly sample per video.
    sr : int
        Target audio sample rate in Hz.
    duration : float
        Clip duration in seconds (audio is truncated/padded to sr*duration).
    transform : callable, optional
        Override the default visual frame transform.  If *None*, the standard
        ImageNet-normalised train/val transform is used based on *split*.
    cache_dir : str | Path, optional
        If provided, preprocessed tensors are saved to ``<cache_dir>/<split>/``
        as ``.pt`` files and re-used on subsequent runs to avoid re-processing.
    """

    LABELS: Dict[str, int] = {"real": 0, "fake": 1}

    def __init__(
        self,
        root_dir: str | Path,
        split: str = "train",
        num_frames: int = 125,
        sr: int = 16_000,
        duration: float = 5.0,
        transform: Optional[Callable] = None,
        cache_dir: Optional[str | Path] = None,
    ) -> None:
        self.root_dir   = Path(root_dir)
        self.split      = split
        self.num_frames = num_frames
        self.sr         = sr
        self.duration   = duration
        self.transform  = transform or (
            VISUAL_TRAIN_TRANSFORM if split == "train" else VISUAL_VAL_TRANSFORM
        )
        self.cache_dir  = Path(cache_dir) / split if cache_dir else None

        if self.cache_dir:
            self.cache_dir.mkdir(parents=True, exist_ok=True)

        # Gather samples as (path, label)
        self.samples: List[Tuple[Path, int]] = []
        for cls_name, label in self.LABELS.items():
            cls_dir = self.root_dir / cls_name
            if not cls_dir.is_dir():
                continue
            for vid in _scan_videos(cls_dir):
                self.samples.append((vid, label))

        if not self.samples:
            raise FileNotFoundError(
                f"No video files found under {self.root_dir}. "
                "Expected sub-folders named 'real' and 'fake'."
            )

        # Shuffle deterministically (same seed for same split)
        rng = random.Random({"train": 42, "val": 43, "test": 44}.get(split, 42))
        rng.shuffle(self.samples)

    # ------------------------------------------------------------------
    # Dataset interface
    # ------------------------------------------------------------------

    def __len__(self) -> int:
        return len(self.samples)

    def __getitem__(self, idx: int) -> Dict[str, object]:
        video_path, label = self.samples[idx]

        # ── Caching ──────────────────────────────────────────────────
        if self.cache_dir is not None:
            cache_file = self.cache_dir / f"{video_path.stem}_{idx}.pt"
            if cache_file.exists():
                try:
                    sample = torch.load(cache_file, map_location="cpu", weights_only=True)
                    sample["video_path"] = str(video_path)
                    sample["label"]      = label
                    return sample
                except Exception:
                    # Corrupt cache — re-process and overwrite
                    pass

        # ── Preprocessing ─────────────────────────────────────────────
        sample = self._preprocess(video_path, label)

        if self.cache_dir is not None:
            # Save without label/path keys so the cache is label-agnostic
            payload = {k: v for k, v in sample.items()
                       if k not in ("label", "video_path")}
            torch.save(payload, cache_file)

        return sample

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _preprocess(self, video_path: Path, label: int) -> Dict[str, object]:
        """Run the full preprocessing pipeline for a single video."""
        frames   = extract_frames(str(video_path), self.num_frames, self.transform)
        waveform = extract_audio_waveform(str(video_path), self.sr, self.duration)
        mel_specs = extract_mel_spectrogram(
            waveform, self.sr, num_frames=self.num_frames
        )
        mouth_crops = extract_mouth_crops(str(video_path), self.num_frames)

        return {
            "frames":      frames,       # (T, 3, 224, 224)
            "waveform":    waveform,     # (sr*duration,)
            "mouth_crops": mouth_crops,  # (T, 3, 96, 96)
            "mel_specs":   mel_specs,    # (T, 1, 80, W)
            "label":       label,
            "video_path":  str(video_path),
        }

    # ------------------------------------------------------------------
    # Class-level utilities
    # ------------------------------------------------------------------

    @classmethod
    def create_splits(
        cls,
        src_dir: str | Path,
        dst_dir: str | Path,
        train_ratio: float = 0.70,
        val_ratio:   float = 0.15,
        seed: int = 42,
        copy: bool = False,
    ) -> Dict[str, "MultiModalDataset"]:
        """
        Split a flat folder of real/fake videos into train/val/test sub-splits.

        The source directory must contain 'real' and 'fake' sub-folders.
        Creates::

            dst_dir/
              train/real/, train/fake/
              val/real/,   val/fake/
              test/real/,  test/fake/

        Parameters
        ----------
        src_dir : path-like
            Source directory with 'real' and 'fake' folders.
        dst_dir : path-like
            Destination for the split folders.
        train_ratio : float
            Fraction of data to use for training.
        val_ratio : float
            Fraction for validation (remainder → test).
        seed : int
            Random seed.
        copy : bool
            If *True*, files are **copied** (not symlinked / moved).
            Use on Windows where symlinks may require elevated permissions.

        Returns
        -------
        dict
            ``{'train': MultiModalDataset, 'val': ..., 'test': ...}``
        """
        src_dir = Path(src_dir)
        dst_dir = Path(dst_dir)
        rng     = random.Random(seed)

        for cls_name in ("real", "fake"):
            videos = _scan_videos(src_dir / cls_name)
            rng.shuffle(videos)

            n_total = len(videos)
            n_train = int(n_total * train_ratio)
            n_val   = int(n_total * val_ratio)

            splits = {
                "train": videos[:n_train],
                "val":   videos[n_train : n_train + n_val],
                "test":  videos[n_train + n_val :],
            }

            for split_name, vids in splits.items():
                split_cls_dir = dst_dir / split_name / cls_name
                split_cls_dir.mkdir(parents=True, exist_ok=True)
                for v in vids:
                    dst = split_cls_dir / v.name
                    if not dst.exists():
                        if copy:
                            shutil.copy2(v, dst)
                        else:
                            # Use relative symlink when possible; fall back to copy
                            try:
                                os.symlink(v.resolve(), dst)
                            except (OSError, NotImplementedError):
                                shutil.copy2(v, dst)

        return {
            split: cls(dst_dir / split, split=split)
            for split in ("train", "val", "test")
        }

    def class_counts(self) -> Dict[str, int]:
        """Return {class_name: count} for logging / weighted sampling."""
        real = sum(1 for _, lbl in self.samples if lbl == 0)
        fake = sum(1 for _, lbl in self.samples if lbl == 1)
        return {"real": real, "fake": fake}

    def __repr__(self) -> str:
        counts = self.class_counts()
        return (
            f"MultiModalDataset(split={self.split!r}, "
            f"total={len(self)}, real={counts['real']}, fake={counts['fake']})"
        )


# ---------------------------------------------------------------------------
# DFDCDataset
# ---------------------------------------------------------------------------

class DFDCDataset(Dataset):
    """
    DFDC Preview dataset loader.

    Expects the DFDC folder structure produced by Kaggle download::

        dfdc_dir/
          metadata.json          ← {filename: {"label": "FAKE"|"REAL", ...}, ...}
          aaa.mp4
          bbb.mp4
          ...

    The DFDC Preview dataset may also be organised in part-folders (00–49).
    This class handles both flat and nested layouts by searching recursively
    for ``metadata.json`` files and their sibling videos.

    Parameters
    ----------
    dfdc_dir : str | Path
        Root directory of (part of) the DFDC Preview dataset.
    split : str
        One of ``{'train', 'val', 'test'}``.  A deterministic 70/15/15
        split is computed from the video list when no external split
        configuration is supplied.
    split_indices : list[int], optional
        If provided, only these integer indices into the full sample list are
        used.  Useful for passing pre-defined train/val/test index lists.
    num_frames : int
        Number of frames to uniformly sample per video.
    sr : int
        Target audio sample rate in Hz.
    duration : float
        Clip duration in seconds.
    transform : callable, optional
        Override the default visual frame transform.
    cache_dir : str | Path, optional
        Directory to cache preprocessed ``.pt`` tensors.
    """

    def __init__(
        self,
        dfdc_dir: str | Path,
        split: str = "train",
        split_indices: Optional[List[int]] = None,
        num_frames: int = 125,
        sr: int = 16_000,
        duration: float = 5.0,
        transform: Optional[Callable] = None,
        cache_dir: Optional[str | Path] = None,
    ) -> None:
        self.dfdc_dir   = Path(dfdc_dir)
        self.split      = split
        self.num_frames = num_frames
        self.sr         = sr
        self.duration   = duration
        self.transform  = transform or (
            VISUAL_TRAIN_TRANSFORM if split == "train" else VISUAL_VAL_TRANSFORM
        )
        self.cache_dir  = Path(cache_dir) / split if cache_dir else None

        if self.cache_dir:
            self.cache_dir.mkdir(parents=True, exist_ok=True)

        # ── Parse metadata.json(s) ────────────────────────────────────
        all_samples: List[Tuple[Path, int]] = []
        meta_files  = list(self.dfdc_dir.rglob("metadata.json"))

        if not meta_files:
            raise FileNotFoundError(
                f"No metadata.json found under {self.dfdc_dir}."
            )

        for meta_file in sorted(meta_files):
            part_dir = meta_file.parent
            with open(meta_file, "r", encoding="utf-8") as fh:
                meta: Dict[str, dict] = json.load(fh)

            for filename, info in meta.items():
                vid_path = part_dir / filename
                if not vid_path.exists():
                    continue

                raw_label = info.get("label", "REAL")
                # DFDC uses "FAKE"/"REAL" strings
                label = 1 if str(raw_label).upper() == "FAKE" else 0
                all_samples.append((vid_path, label))

        if not all_samples:
            raise RuntimeError(
                f"metadata.json found but no matching video files located under "
                f"{self.dfdc_dir}. Check your dataset download."
            )

        # ── Split ──────────────────────────────────────────────────────
        if split_indices is not None:
            self.samples = [all_samples[i] for i in split_indices]
        else:
            self.samples = self._auto_split(all_samples, split)

    # ------------------------------------------------------------------
    # Dataset interface
    # ------------------------------------------------------------------

    def __len__(self) -> int:
        return len(self.samples)

    def __getitem__(self, idx: int) -> Dict[str, object]:
        video_path, label = self.samples[idx]

        # ── Caching ──────────────────────────────────────────────────
        if self.cache_dir is not None:
            safe_stem   = video_path.stem.replace("/", "_")
            cache_file  = self.cache_dir / f"{safe_stem}_{idx}.pt"
            if cache_file.exists():
                try:
                    sample = torch.load(cache_file, map_location="cpu", weights_only=True)
                    sample["video_path"] = str(video_path)
                    sample["label"]      = label
                    return sample
                except Exception:
                    pass  # corrupt cache; re-process

        # ── Preprocessing ─────────────────────────────────────────────
        sample = self._preprocess(video_path, label)

        if self.cache_dir is not None:
            payload = {k: v for k, v in sample.items()
                       if k not in ("label", "video_path")}
            torch.save(payload, cache_file)

        return sample

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _preprocess(self, video_path: Path, label: int) -> Dict[str, object]:
        """Full preprocessing pipeline for one DFDC video."""
        frames      = extract_frames(str(video_path), self.num_frames, self.transform)
        waveform    = extract_audio_waveform(str(video_path), self.sr, self.duration)
        mel_specs   = extract_mel_spectrogram(
            waveform, self.sr, num_frames=self.num_frames
        )
        mouth_crops = extract_mouth_crops(str(video_path), self.num_frames)

        return {
            "frames":      frames,
            "waveform":    waveform,
            "mouth_crops": mouth_crops,
            "mel_specs":   mel_specs,
            "label":       label,
            "video_path":  str(video_path),
        }

    @staticmethod
    def _auto_split(
        samples: List[Tuple[Path, int]],
        split: str,
        train_ratio: float = 0.70,
        val_ratio:   float = 0.15,
        seed: int = 42,
    ) -> List[Tuple[Path, int]]:
        """Deterministically split *samples* into train/val/test."""
        rng = random.Random(seed)
        samples = list(samples)
        rng.shuffle(samples)

        n       = len(samples)
        n_train = int(n * train_ratio)
        n_val   = int(n * val_ratio)

        if split == "train":
            return samples[:n_train]
        elif split == "val":
            return samples[n_train : n_train + n_val]
        elif split == "test":
            return samples[n_train + n_val :]
        else:
            raise ValueError(f"Unknown split: {split!r}. Choose 'train', 'val', or 'test'.")

    def class_counts(self) -> Dict[str, int]:
        """Return {class_name: count}."""
        real = sum(1 for _, lbl in self.samples if lbl == 0)
        fake = sum(1 for _, lbl in self.samples if lbl == 1)
        return {"real": real, "fake": fake}

    def get_weighted_sampler(self) -> "torch.utils.data.WeightedRandomSampler":
        """
        Build a WeightedRandomSampler to address class imbalance.

        Usage::

            sampler = dataset.get_weighted_sampler()
            loader  = DataLoader(dataset, batch_size=8, sampler=sampler)
        """
        from torch.utils.data import WeightedRandomSampler

        counts  = self.class_counts()
        n_real, n_fake = counts["real"], counts["fake"]
        n_total = n_real + n_fake

        w_real = n_total / (2.0 * n_real) if n_real > 0 else 0.0
        w_fake = n_total / (2.0 * n_fake) if n_fake > 0 else 0.0

        weights = [
            w_real if lbl == 0 else w_fake
            for _, lbl in self.samples
        ]
        return WeightedRandomSampler(
            weights=weights,
            num_samples=len(weights),
            replacement=True,
        )

    def __repr__(self) -> str:
        counts = self.class_counts()
        return (
            f"DFDCDataset(split={self.split!r}, "
            f"total={len(self)}, real={counts['real']}, fake={counts['fake']}, "
            f"root={self.dfdc_dir})"
        )


# ---------------------------------------------------------------------------
# Quick self-test
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import sys

    if len(sys.argv) < 2:
        print("Usage: python datasets.py <path_to_dataset_root> [dfdc|multimodal]")
        print("  dfdc       → expects metadata.json in the root")
        print("  multimodal → expects real/ and fake/ sub-folders")
        sys.exit(0)

    root = sys.argv[1]
    mode = sys.argv[2] if len(sys.argv) > 2 else "multimodal"

    if mode == "dfdc":
        ds = DFDCDataset(root, split="train", num_frames=8)
    else:
        ds = MultiModalDataset(root, split="train", num_frames=8)

    print(ds)
    sample = ds[0]
    print(f"  frames      : {sample['frames'].shape}")
    print(f"  waveform    : {sample['waveform'].shape}")
    print(f"  mouth_crops : {sample['mouth_crops'].shape}")
    print(f"  mel_specs   : {sample['mel_specs'].shape}")
    print(f"  label       : {sample['label']}")
    print(f"  video_path  : {sample['video_path']}")