Blame for src/nida/train/data.py

"""Latent dataset for training: on-disk format, loader and rejection report.

On-disk format (a *latent dataset directory*)::

    <dir>/index.jsonl          one JSON object per segment (UTF-8, one per line)
    <dir>/shard-00000.npy      float16 or float32 array [total_frames, latent_dim]
    <dir>/shard-00001.npy      ...

Each ``index.jsonl`` record has exactly these keys:

==================  ==========================================================================
``segment_id``      non-empty string, unique across the index
``shard``           file name of a shard in the same directory (no path separators)
``offset``          int >= 0, first frame row of the segment inside the shard
``num_frames``      int >= 1, number of latent frames; ``offset + num_frames <= rows``
``text``            tag-spec tagged Turkish text (parsed by :func:`nida.model.vocab.encode_text`)
``emotion``         one of :data:`EMOTIONS` or ``null``
``intensity``       float in [0, 1] or ``null``
``rate_bucket``     int in 0..4 or ``null``
``loudness_bucket`` int in 0..2 or ``null``
``group``           non-empty string (speaker / source video id; splits never cross groups)
``split``           ``train`` | ``val`` | ``test``
==================  ==========================================================================

Optional key: ``language`` — one of :data:`LANGUAGES` (ADR-0007); ``"tr"`` when absent, so
datasets written before it stay valid.

``null`` condition values map to the model's *null* embedding (the id one past the last real
class). A segment of ``n`` frames is the latent slice ``shard[offset : offset + n]``; all
shards share one ``latent_dim``. Shards are opened with ``np.load(mmap_mode="r")`` and
converted to float32 only for the rows of a requested segment.

Structural problems (bad JSON, missing/extra keys, wrong types, duplicate ids, missing
shard, out-of-range rows, inconsistent ``latent_dim``, a ``group`` in more than one
split) raise :class:`DataFormatError` --
they mean the producer is broken. *Item-level* problems that are properties of the data are
**rejected with a counted reason** and reported in :class:`LoadReport` (never dropped
silently): ``text_longer_than_frames`` (the F5-style text padding needs
``len(text_ids) <= num_frames``), ``text_encoding_error`` (unknown character / bad tag) and
``empty_text``.

Several directories (one per data-pilot root) can be trained on as one union with
:func:`load_latent_datasets`; see its docstring for the cross-directory checks.
"""

from __future__ import annotations

import hashlib
import json
import re
from collections import Counter
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field, replace
from pathlib import Path

import numpy as np
import numpy.typing as npt

from nida.data.timings import SegmentTimings, load_timings
from nida.model.vocab import encode_text

__all__ = [
    "EMOTIONS",
    "LANGUAGES",
    "N_LOUDNESS_BUCKETS",
    "N_RATE_BUCKETS",
    "REJECT_EMPTY_TEXT",
    "REJECT_TEXT_ENCODING",
    "REJECT_TEXT_TOO_LONG",
    "SPLITS",
    "DataFormatError",
    "LatentDataset",
    "LatentStats",
    "LoadReport",
    "Rejection",
    "Segment",
    "compute_latent_stats",
    "load_latent_dataset",
    "load_latent_datasets",
]

#: Language order (ADR-0007); the id is the index and must match ``ModelConfig.n_languages``
#: in models trained with languages. Append-only.
LANGUAGES: tuple[str, ...] = ("tr", "en")

#: Fixed emotion order; the id of an emotion is its index, the null id is ``len(EMOTIONS)``
#: (== ``ModelConfig.n_emotions``). Append-only: changing it invalidates trained models.
EMOTIONS: tuple[str, ...] = (
    "angry",
    "disgusted",
    "fearful",
    "happy",
    "neutral",
    "sad",
    "surprised",
)
N_RATE_BUCKETS = 5
N_LOUDNESS_BUCKETS = 3
SPLITS: tuple[str, ...] = ("train", "val", "test")

REJECT_TEXT_TOO_LONG = "text_longer_than_frames"
REJECT_TEXT_ENCODING = "text_encoding_error"
REJECT_EMPTY_TEXT = "empty_text"
REJECT_NO_TIMINGS = "no_timings"  # E11 aligned text path: segment missing from timings.jsonl
REJECT_TIMINGS_MISMATCH = "timings_mismatch"  # token / frame count differs from the index

INDEX_NAME = "index.jsonl"
_SHARD_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*\.npy$")
_RECORD_KEYS = (
    "segment_id",
    "shard",
    "offset",
    "num_frames",
    "text",
    "emotion",
    "intensity",
    "rate_bucket",
    "loudness_bucket",
    "group",
    "split",
)

Float32Array = npt.NDArray[np.float32]


class DataFormatError(ValueError):
    """The latent dataset on disk violates the documented format."""


@dataclass(frozen=True)
class Segment:
    """One training segment. ``text_ids`` are filled in by the loader (empty before)."""

    segment_id: str
    shard: str
    offset: int
    num_frames: int
    text: str
    emotion: str | None
    intensity: float | None
    rate_bucket: int | None
    loudness_bucket: int | None
    group: str
    split: str
    text_ids: tuple[int, ...] = ()
    language: str = "tr"
    #: E11 timeline (aligned text path): ``(char_start, char_end, frame_start)`` per unit,
    #: tiling ``text_ids`` and the frames (``nida.data.timings``); empty when not loaded.
    units: tuple[tuple[int, int, int], ...] = ()


@dataclass(frozen=True)
class Rejection:
    """An item that was loaded from the index but not admitted to the dataset."""

    segment_id: str
    reason: str
    detail: str


@dataclass
class LoadReport:
    """What the loader did with the index (every record is accounted for)."""

    data_dir: str
    split: str
    num_index_records: int = 0
    num_in_split: int = 0
    num_accepted: int = 0
    accepted_frames: int = 0
    latent_dim: int = 0
    rejection_counts: dict[str, int] = field(default_factory=dict)
    rejections: list[Rejection] = field(default_factory=list)
    #: Per-directory reports of a merged multi-directory load (empty for a single directory).
    per_dir: list[LoadReport] = field(default_factory=list)

    @property
    def num_rejected(self) -> int:
        return sum(self.rejection_counts.values())

    def to_dict(self, *, max_examples: int = 20) -> dict[str, object]:
        """JSON summary (the first ``max_examples`` rejections as examples).

        A merged multi-directory report adds a ``per_dir`` list (one summary per directory,
        in load order) and joins the directories with ``"; "`` in ``data_dir``; the key is
        absent for a single directory, whose summary is unchanged.
        """
        summary: dict[str, object] = {
            "data_dir": self.data_dir,
            "split": self.split,
            "num_index_records": self.num_index_records,
            "num_in_split": self.num_in_split,
            "num_accepted": self.num_accepted,
            "num_rejected": self.num_rejected,
            "accepted_frames": self.accepted_frames,
            "latent_dim": self.latent_dim,
            "rejection_counts": dict(sorted(self.rejection_counts.items())),
            "rejection_examples": [
                {"segment_id": r.segment_id, "reason": r.reason, "detail": r.detail}
                for r in self.rejections[:max_examples]
            ],
        }
        if self.per_dir:
            summary["per_dir"] = [r.to_dict(max_examples=max_examples) for r in self.per_dir]
        return summary


def _fail(line_no: int, message: str) -> DataFormatError:
    return DataFormatError(f"{INDEX_NAME} line {line_no}: {message}")


def _is_int(value: object) -> bool:
    return isinstance(value, int) and not isinstance(value, bool)


def _parse_record(raw: object, line_no: int) -> Segment:
    """Validate one decoded JSON object and build a :class:`Segment` (no text ids yet)."""
    if not isinstance(raw, dict):
        raise _fail(line_no, f"expected a JSON object, got {type(raw).__name__}")
    keys = set(raw) - {"language"}  # optional (ADR-0007), "tr" when absent
    if keys != set(_RECORD_KEYS):
        raise _fail(
            line_no,
            f"keys differ: missing {sorted(set(_RECORD_KEYS) - keys)}, "
            f"unknown {sorted(keys - set(_RECORD_KEYS))}",
        )
    segment_id, shard, text, group, split = (
        raw["segment_id"],
        raw["shard"],
        raw["text"],
        raw["group"],
        raw["split"],
    )
    for name, value in (
        ("segment_id", segment_id),
        ("shard", shard),
        ("text", text),
        ("group", group),
    ):
        if not isinstance(value, str) or not value:
            raise _fail(line_no, f"{name} must be a non-empty string, got {value!r}")
    if not _SHARD_RE.fullmatch(shard):
        raise _fail(line_no, f"shard must be a plain *.npy file name, got {shard!r}")
    if split not in SPLITS:
        raise _fail(line_no, f"split must be one of {SPLITS}, got {split!r}")
    offset, num_frames = raw["offset"], raw["num_frames"]
    if not _is_int(offset) or offset < 0:
        raise _fail(line_no, f"offset must be an int >= 0, got {offset!r}")
    if not _is_int(num_frames) or num_frames < 1:
        raise _fail(line_no, f"num_frames must be an int >= 1, got {num_frames!r}")
    emotion = raw["emotion"]
    if emotion is not None and emotion not in EMOTIONS:
        raise _fail(line_no, f"emotion must be null or one of {EMOTIONS}, got {emotion!r}")
    intensity = raw["intensity"]
    if intensity is not None:
        if isinstance(intensity, bool) or not isinstance(intensity, int | float):
            raise _fail(line_no, f"intensity must be null or a number, got {intensity!r}")
        intensity = float(intensity)
        if not 0.0 <= intensity <= 1.0:  # also rejects NaN
            raise _fail(line_no, f"intensity must be in [0, 1], got {intensity!r}")
    for name, upper in (("rate_bucket", N_RATE_BUCKETS), ("loudness_bucket", N_LOUDNESS_BUCKETS)):
        value = raw[name]
        if value is not None and (not _is_int(value) or not 0 <= value < upper):
            raise _fail(line_no, f"{name} must be null or an int in 0..{upper - 1}, got {value!r}")
    language = raw.get("language", "tr")
    if language not in LANGUAGES:
        raise _fail(line_no, f"language must be one of {LANGUAGES}, got {language!r}")
    return Segment(
        segment_id=segment_id,
        shard=shard,
        offset=offset,
        num_frames=num_frames,
        text=text,
        emotion=emotion,
        intensity=intensity,
        rate_bucket=raw["rate_bucket"],
        loudness_bucket=raw["loudness_bucket"],
        group=group,
        split=split,
        language=language,
    )


def _read_index(path: Path) -> list[Segment]:
    if not path.is_file():
        raise DataFormatError(f"missing {path}")
    records: list[Segment] = []
    seen: set[str] = set()
    with path.open(encoding="utf-8") as handle:
        for line_no, line in enumerate(handle, start=1):
            if not line.strip():
                continue
            try:
                raw = json.loads(line)
            except json.JSONDecodeError as exc:
                raise _fail(line_no, f"invalid JSON ({exc})") from exc
            record = _parse_record(raw, line_no)
            if record.segment_id in seen:
                raise _fail(line_no, f"duplicate segment_id {record.segment_id!r}")
            seen.add(record.segment_id)
            records.append(record)
    return records


class LatentDataset:
    """Accepted segments of one split, with memory-mapped latent shards.

    ``data_dir`` is the (first) directory; ``data_dirs`` lists every directory of a merged
    multi-directory dataset (see :func:`load_latent_datasets`), whose ``Segment.shard`` and
    shard keys are then ``"<dir index>/<file name>"``.

    Not thread-safe to :meth:`close` while in use. ``close()`` drops the memory maps (needed
    on Windows before deleting the files).
    """

    def __init__(
        self,
        data_dir: Path,
        segments: Sequence[Segment],
        shards: dict[str, npt.NDArray[np.floating]],
        latent_dim: int,
        *,
        data_dirs: Sequence[Path] | None = None,
    ) -> None:
        self.data_dir = data_dir
        self.data_dirs: tuple[Path, ...] = (data_dir,) if data_dirs is None else tuple(data_dirs)
        self.segments: tuple[Segment, ...] = tuple(segments)
        self._shards = shards
        self.latent_dim = latent_dim

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

    def __getitem__(self, index: int) -> Segment:
        return self.segments[index]

    @property
    def lengths(self) -> list[int]:
        """``num_frames`` per segment, in dataset order."""
        return [s.num_frames for s in self.segments]

    def load_latents(self, index: int) -> Float32Array:
        """The segment's latent ``[num_frames, latent_dim]`` as a float32 copy.

        Raises:
            DataFormatError: the stored values contain NaN/Inf.
        """
        seg = self.segments[index]
        rows = self._shards[seg.shard][seg.offset : seg.offset + seg.num_frames]
        out: Float32Array = np.array(rows, dtype=np.float32)
        if not bool(np.isfinite(out).all()):
            raise DataFormatError(f"non-finite latent values in segment {seg.segment_id!r}")
        return out

    def fingerprint(self) -> str:
        """SHA-256 over (id, shard, offset, frames) of every accepted segment, in order.

        Stored in checkpoints: a resume against different data (or a different rejection
        outcome) must fail, since the sampler position would be meaningless.
        """
        digest = hashlib.sha256()
        for s in self.segments:
            digest.update(f"{s.segment_id}\t{s.shard}\t{s.offset}\t{s.num_frames}\n".encode())
            if s.units:  # E11 timeline: a changed timings.jsonl must not resume silently
                digest.update(f"units\t{s.units}\n".encode())
        return digest.hexdigest()

    def close(self) -> None:
        self._shards.clear()


#: Channels whose spread is below this are not scaled up (avoids amplifying noise).
MIN_LATENT_STD = 1e-3


@dataclass(frozen=True)
class LatentStats:
    """Per-channel mean / std of the training latents (the model sees ``(x - mean) / std``).

    Flow matching draws ``x0 ~ N(0, I)``; raw VAE latents have arbitrary per-channel offset
    and scale, so without this the noise and data ends of the path are mismatched. The
    stats belong to the model: they are stored in every checkpoint and are needed to turn
    generated latents back into VAE latents (:meth:`denormalize`).
    """

    mean: tuple[float, ...]
    std: tuple[float, ...]

    def __post_init__(self) -> None:
        if not self.mean or len(self.mean) != len(self.std):
            raise ValueError("mean and std must be non-empty and of equal length")
        if not all(np.isfinite(self.mean)) or not all(
            np.isfinite(s) and s >= MIN_LATENT_STD for s in self.std
        ):
            raise ValueError(f"latent stats must be finite with std >= {MIN_LATENT_STD}")

    def close_to(self, other: LatentStats, rtol: float = 1e-6) -> bool:
        """Same channels and values equal up to float summation noise."""
        return len(self.mean) == len(other.mean) and bool(
            np.allclose(self.mean, other.mean, rtol=rtol, atol=1e-9)
            and np.allclose(self.std, other.std, rtol=rtol, atol=1e-9)
        )

    def normalize(self, latents: Float32Array) -> Float32Array:
        """``latents: (num_frames, latent_dim)`` -> standardised float32 copy."""
        mean = np.asarray(self.mean, dtype=np.float32)
        std = np.asarray(self.std, dtype=np.float32)
        out: Float32Array = ((latents - mean) / std).astype(np.float32)
        return out

    def denormalize(self, latents: Float32Array) -> Float32Array:
        """Inverse of :meth:`normalize`."""
        mean = np.asarray(self.mean, dtype=np.float32)
        std = np.asarray(self.std, dtype=np.float32)
        out: Float32Array = (latents * std + mean).astype(np.float32)
        return out

    def to_dict(self) -> dict[str, list[float]]:
        return {"mean": list(self.mean), "std": list(self.std)}

    @classmethod
    def from_mapping(cls, raw: object) -> LatentStats:
        if not isinstance(raw, dict) or set(raw) != {"mean", "std"}:
            raise ValueError(f"latent stats must be a mapping with mean/std, got {raw!r}")
        return cls(tuple(float(x) for x in raw["mean"]), tuple(float(x) for x in raw["std"]))


def compute_latent_stats(data: LatentDataset) -> LatentStats:
    """Per-channel mean/std over every frame of ``data`` (float64 accumulation).

    ``std`` is floored at :data:`MIN_LATENT_STD`.

    Raises:
        ValueError: ``data`` is empty.
    """
    if len(data) == 0:
        raise ValueError("cannot compute latent stats of an empty dataset")
    total = np.zeros(data.latent_dim, dtype=np.float64)
    total_sq = np.zeros(data.latent_dim, dtype=np.float64)
    count = 0
    for index in range(len(data)):
        lat = data.load_latents(index).astype(np.float64)
        total += lat.sum(axis=0)
        total_sq += np.square(lat).sum(axis=0)
        count += lat.shape[0]
    mean = total / count
    var = np.maximum(total_sq / count - np.square(mean), 0.0)
    std = np.maximum(np.sqrt(var), MIN_LATENT_STD)
    return LatentStats(tuple(float(x) for x in mean), tuple(float(x) for x in std))


def _open_shards(
    data_dir: Path, records: Sequence[Segment]
) -> tuple[dict[str, npt.NDArray[np.floating]], int]:
    """Memory-map every referenced shard and validate bounds. Returns (shards, latent_dim)."""
    shards: dict[str, npt.NDArray[np.floating]] = {}
    latent_dim: int | None = None
    for name in sorted({r.shard for r in records}):
        path = data_dir / name
        if not path.is_file():
            raise DataFormatError(f"shard {name!r} referenced by the index is missing")
        try:
            arr = np.load(path, mmap_mode="r", allow_pickle=False)
        except (ValueError, OSError) as exc:
            raise DataFormatError(f"cannot read shard {name!r}: {exc}") from exc
        if arr.ndim != 2 or arr.dtype not in (np.float16, np.float32):
            raise DataFormatError(
                f"shard {name!r} must be a 2-D float16/float32 array, "
                f"got dtype {arr.dtype} shape {arr.shape}"
            )
        if latent_dim is None:
            latent_dim = int(arr.shape[1])
        elif int(arr.shape[1]) != latent_dim:
            raise DataFormatError(
                f"shard {name!r} has latent_dim {arr.shape[1]}, expected {latent_dim}"
            )
        shards[name] = arr
    for r in records:
        rows = shards[r.shard].shape[0]
        if r.offset + r.num_frames > rows:
            raise DataFormatError(
                f"segment {r.segment_id!r}: rows {r.offset}..{r.offset + r.num_frames} "
                f"exceed shard {r.shard!r} ({rows} rows)"
            )
    return shards, latent_dim if latent_dim is not None else 0


def _encode(segment: Segment) -> tuple[Segment | None, Rejection | None]:
    """Encode the text; return the segment with ids or the rejection."""
    try:
        ids = encode_text(segment.text)
    except ValueError as exc:  # TagSyntaxError subclasses ValueError; vocab raises ValueError
        return None, Rejection(segment.segment_id, REJECT_TEXT_ENCODING, str(exc))
    if not ids:
        return None, Rejection(segment.segment_id, REJECT_EMPTY_TEXT, "text encodes to no tokens")
    if len(ids) > segment.num_frames:
        detail = f"{len(ids)} text tokens > {segment.num_frames} frames"
        return None, Rejection(segment.segment_id, REJECT_TEXT_TOO_LONG, detail)
    return replace(segment, text_ids=tuple(ids)), None


def _check_groups_disjoint(records: Sequence[Segment]) -> None:
    """A group (speaker / source video) in two splits would leak voices into evaluation."""
    splits_of: dict[str, set[str]] = {}
    for record in records:
        splits_of.setdefault(record.group, set()).add(record.split)
    crossing = sorted(g for g, s in splits_of.items() if len(s) > 1)
    if crossing:
        raise DataFormatError(
            f"{len(crossing)} group(s) appear in more than one split, e.g. "
            f"{crossing[0]!r} in {sorted(splits_of[crossing[0]])}"
        )


def _attach_timings(
    segment: Segment, timings: Mapping[str, SegmentTimings]
) -> tuple[Segment | None, Rejection | None]:
    """E11: attach the unit timeline; reject a segment without one or with a mismatch."""
    found = timings.get(segment.segment_id)
    if found is None:
        return None, Rejection(segment.segment_id, REJECT_NO_TIMINGS, "not in timings.jsonl")
    if found.n_tokens != len(segment.text_ids) or found.num_frames != segment.num_frames:
        detail = (
            f"timings {found.n_tokens} tokens / {found.num_frames} frames != "
            f"{len(segment.text_ids)} tokens / {segment.num_frames} frames"
        )
        return None, Rejection(segment.segment_id, REJECT_TIMINGS_MISMATCH, detail)
    units = tuple((u.char_start, u.char_end, u.frame_start) for u in found.units)
    return replace(segment, units=units), None


def _load_split(
    data_dir: Path, split: str, records: Sequence[Segment], *, timings: bool = False
) -> tuple[LatentDataset, LoadReport]:
    """Open, validate and encode the ``split`` records of one directory's (checked) index.

    With ``timings`` (aligned text path) every accepted segment also gets its unit
    timeline from ``timings.jsonl`` (:func:`nida.data.timings.load_timings`).
    """
    unit_map = load_timings(data_dir) if timings else None
    in_split = [r for r in records if r.split == split]
    shards, latent_dim = _open_shards(data_dir, in_split)
    report = LoadReport(
        data_dir=str(data_dir),
        split=split,
        num_index_records=len(records),
        num_in_split=len(in_split),
        latent_dim=latent_dim,
    )
    accepted: list[Segment] = []
    counts: Counter[str] = Counter()
    for record in in_split:
        encoded, rejection = _encode(record)
        if encoded is not None and unit_map is not None:
            encoded, rejection = _attach_timings(encoded, unit_map)
        if encoded is not None:
            accepted.append(encoded)
        elif rejection is not None:
            counts[rejection.reason] += 1
            report.rejections.append(rejection)
    report.rejection_counts = dict(counts)
    report.num_accepted = len(accepted)
    report.accepted_frames = sum(s.num_frames for s in accepted)
    if report.num_accepted + report.num_rejected != report.num_in_split:
        raise AssertionError("loader lost track of records")  # invariant, not input error
    return LatentDataset(data_dir, accepted, shards, latent_dim), report


def load_latent_dataset(
    data_dir: Path, split: str, *, timings: bool = False
) -> tuple[LatentDataset, LoadReport]:
    """Load one split of a latent dataset directory (format: see module docstring).

    Returns the dataset (accepted items, in index order) and the report that accounts for
    every record of the split. The caller decides whether an empty dataset is fatal.

    Raises:
        DataFormatError: structural problems (see module docstring).
        ValueError: unknown ``split``.
    """
    if split not in SPLITS:
        raise ValueError(f"split must be one of {SPLITS}, got {split!r}")
    records = _read_index(data_dir / INDEX_NAME)
    _check_groups_disjoint(records)
    return _load_split(data_dir, split, records, timings=timings)


def _check_cross_directory(data_dirs: Sequence[Path], per_dir: Sequence[Sequence[Segment]]) -> None:
    """Checks that only make sense across directories (each index is already valid alone)."""
    owner: dict[str, int] = {}
    for k, records in enumerate(per_dir):
        for record in records:
            first = owner.setdefault(record.segment_id, k)
            if first != k:
                raise DataFormatError(
                    f"segment_id {record.segment_id!r} appears in both {data_dirs[first]} "
                    f"and {data_dirs[k]}"
                )
    try:
        _check_groups_disjoint([r for records in per_dir for r in records])
    except DataFormatError as exc:
        raise DataFormatError(
            f"across data directories {[str(d) for d in data_dirs]}: {exc}"
        ) from exc


#: ``meta.json`` keys that must agree before directories may be merged: the condition
#: buckets are only comparable when their edges are equal (the pilot writes fixed edges
#: from configs/data/dataset.yaml), and frames must have the same duration.
META_MUST_MATCH = (
    "frame_hop_ms",
    "rate_bucket_edges_chars_per_s",
    "loudness_bucket_edges_dbfs",
)
META_NAME = "meta.json"


def _check_meta_compatible(data_dirs: Sequence[Path]) -> None:
    """Directories with a ``meta.json`` must agree on :data:`META_MUST_MATCH`; a mix of
    directories with and without one is refused (the missing one cannot be checked)."""
    metas: list[tuple[Path, Mapping[str, object]]] = []
    for data_dir in data_dirs:
        path = data_dir / META_NAME
        if path.is_file():
            try:
                loaded = json.loads(path.read_text(encoding="utf-8"))
            except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc:
                raise DataFormatError(f"{path}: unreadable meta.json: {exc}") from exc
            if not isinstance(loaded, dict):
                raise DataFormatError(f"{path}: meta.json must be an object")
            metas.append((data_dir, loaded))
    if not metas:
        return
    if len(metas) != len(data_dirs):
        missing = [str(d) for d in data_dirs if not (d / META_NAME).is_file()]
        raise DataFormatError(f"cannot merge: no {META_NAME} in {missing} (others have one)")
    first_dir, first = metas[0]
    for data_dir, meta in metas[1:]:
        for key in META_MUST_MATCH:
            if meta.get(key) != first.get(key):
                raise DataFormatError(
                    f"cannot merge: {key} differs ({first_dir}: {first.get(key)!r}, "
                    f"{data_dir}: {meta.get(key)!r})"
                )


def _check_latent_dims(data_dirs: Sequence[Path], reports: Sequence[LoadReport]) -> None:
    known = [(d, r.latent_dim) for d, r in zip(data_dirs, reports, strict=True) if r.latent_dim]
    for data_dir, dim in known[1:]:
        if dim != known[0][1]:
            raise DataFormatError(
                f"latent_dim differs across data directories: {known[0][0]} has {known[0][1]}, "
                f"{data_dir} has {dim}"
            )


def _merge_reports(split: str, reports: Sequence[LoadReport]) -> LoadReport:
    counts: Counter[str] = Counter()
    for report in reports:
        counts.update(report.rejection_counts)
    return LoadReport(
        data_dir="; ".join(r.data_dir for r in reports),
        split=split,
        num_index_records=sum(r.num_index_records for r in reports),
        num_in_split=sum(r.num_in_split for r in reports),
        num_accepted=sum(r.num_accepted for r in reports),
        accepted_frames=sum(r.accepted_frames for r in reports),
        latent_dim=max(r.latent_dim for r in reports),
        rejection_counts=dict(counts),
        rejections=[x for r in reports for x in r.rejections],
        per_dir=list(reports),
    )


def load_latent_datasets(
    data_dirs: Sequence[Path], split: str, *, timings: bool = False
) -> tuple[LatentDataset, LoadReport]:
    """Load one split of several latent dataset directories as one union dataset.

    Segments keep directory order, then index order within a directory. With a single
    directory this is exactly :func:`load_latent_dataset` (same segments, same
    ``fingerprint()``, so existing checkpoints still resume). With several, every shard key
    and ``Segment.shard`` is prefixed with ``"<dir index>/"`` to keep shards of different
    directories apart (the fingerprint therefore depends on the directory order), and the
    returned report sums the counts and carries the per-directory reports in ``per_dir``.

    Cross-directory checks (the index records of *all* splits are considered, not just
    ``split``): a ``segment_id`` in two directories, a ``group`` that is in different
    splits in different directories, and differing ``latent_dim``. A directory with no
    records in ``split`` has no shard to read, so its ``latent_dim`` is unknown and is not
    compared.

    Raises:
        DataFormatError: structural problems in any directory, or a failed cross check.
        ValueError: unknown ``split``, no directories, or the same directory given twice.
    """
    if split not in SPLITS:
        raise ValueError(f"split must be one of {SPLITS}, got {split!r}")
    if not data_dirs:
        raise ValueError("at least one latent dataset directory is required")
    if len(data_dirs) == 1:
        return load_latent_dataset(data_dirs[0], split, timings=timings)
    resolved = [d.resolve() for d in data_dirs]
    if len(set(resolved)) != len(resolved):
        raise ValueError(f"the same latent dataset directory is given twice: {list(data_dirs)}")
    per_dir_records: list[list[Segment]] = []
    for data_dir in data_dirs:
        records = _read_index(data_dir / INDEX_NAME)
        _check_groups_disjoint(records)
        per_dir_records.append(records)
    _check_cross_directory(data_dirs, per_dir_records)
    _check_meta_compatible(data_dirs)

    loaded: list[tuple[LatentDataset, LoadReport]] = []
    try:
        for data_dir, records in zip(data_dirs, per_dir_records, strict=True):
            loaded.append(_load_split(data_dir, split, records, timings=timings))
        _check_latent_dims(data_dirs, [report for _, report in loaded])
    except BaseException:
        for dataset, _ in loaded:
            dataset.close()
        raise
    segments: list[Segment] = []
    shards: dict[str, npt.NDArray[np.floating]] = {}
    for k, (dataset, _) in enumerate(loaded):
        segments.extend(replace(s, shard=f"{k}/{s.shard}") for s in dataset.segments)
        shards.update({f"{k}/{name}": arr for name, arr in dataset._shards.items()})
    report = _merge_reports(split, [report for _, report in loaded])
    merged = LatentDataset(
        data_dirs[0], segments, shards, report.latent_dim, data_dirs=list(data_dirs)
    )
    return merged, report