from mapgs.data.types import MapGSSample, collate_samples
from mapgs.data.scene import generate_scene, ProceduralScene, MapNoise
from mapgs.data.synthetic import (
    SyntheticDataset,
    render_scene_gt,
    render_deviated_gt,
    place_dynamic_gt,
)
from mapgs.data.real import (
    RawClip,
    raw_to_sample,
    BaseDrivingDataset,
    WaymoDataset,
    NuScenesDataset,
    Argoverse2Dataset,
)
from mapgs.data.unified import UnifiedClipDataset, write_unified_clip, UNIFIED_VERSION

__all__ = [
    "MapGSSample",
    "collate_samples",
    "generate_scene",
    "ProceduralScene",
    "MapNoise",
    "SyntheticDataset",
    "render_scene_gt",
    "render_deviated_gt",
    "place_dynamic_gt",
    "RawClip",
    "raw_to_sample",
    "BaseDrivingDataset",
    "WaymoDataset",
    "NuScenesDataset",
    "Argoverse2Dataset",
    "UnifiedClipDataset",
    "write_unified_clip",
    "UNIFIED_VERSION",
]


def build_dataset(cfg, split: str = "train", **kwargs):
    """Factory: dispatch on ``cfg.data.name``.

    For ``unified`` (the recommended path for real Waymo+AV2 training), set
    ``cfg.data.root`` to one or more unified roots (comma-separated to train on a
    mix, e.g. ``"data/unified/waymo,data/unified/av2"``).
    """
    name = cfg.data.name.lower()
    if name == "synthetic":
        return SyntheticDataset(cfg, split=split, **kwargs)
    if name == "unified":
        roots = [r.strip() for r in cfg.data.root.split(",") if r.strip()]
        return UnifiedClipDataset(cfg, roots=roots, split=split, **kwargs)
    if name == "waymo":
        return WaymoDataset(cfg, split=split, **kwargs)
    if name == "nuscenes":
        return NuScenesDataset(cfg, split=split, **kwargs)
    if name in ("argoverse", "argoverse2", "av2"):
        return Argoverse2Dataset(cfg, split=split, **kwargs)
    raise ValueError(f"unknown dataset {cfg.data.name!r}")