Blame for src/nida/data/segment.py

"""S3 segmentation: turn VAD speech, speaker turns and overlap regions into segments.

Pure functions over integer-millisecond intervals, all relative to the *source file*
(``file_start_ms`` / ``file_end_ms``; half-open ``[start, end)``). Nothing here reads audio.

Algorithm (see :func:`plan_segments`):

a. *Single-speaker speech regions*: speech, minus the overlap regions expanded by
   ``overlap_guard_ms`` (the "forbidden zone"), restricted to one speaker's turns.
b. Inside a run of one speaker's regions, speech separated by silences shorter than
   ``min_cut_silence_ms`` is merged into one chunk (a pause is not a cut point).
c. Consecutive chunks are packed greedily into segments of at most ``max_segment_ms``;
   a cut falls only in a silence of at least ``min_cut_silence_ms``. A single chunk longer
   than ``max_segment_ms`` is never force-cut; it is rejected.
d. Each segment is padded by ``pad_ms`` per side, but never beyond the file, the middle of
   the neighbouring silence, the forbidden zone, or another speaker's turn.
e. Segments whose speech extent is shorter than ``min_segment_ms`` are rejected.

Nothing is dropped silently: every millisecond of speech ends up either inside a planned
segment or inside a :class:`RejectedSpan` with a machine-readable reason.
"""

from __future__ import annotations

from bisect import bisect_left, bisect_right
from collections.abc import Iterable, Mapping, Sequence
from dataclasses import dataclass
from typing import Literal, get_args

__all__ = [
    "REJECT_REASONS",
    "Interval",
    "PlannedSegment",
    "RejectReason",
    "RejectedSpan",
    "SegmentConfig",
    "SegmentPlan",
    "SpeakerTurn",
    "plan_segments",
]

RejectReason = Literal["too_short", "too_long_without_pause", "overlap", "no_speaker"]
REJECT_REASONS: tuple[str, ...] = get_args(RejectReason)

# Internal representation: half-open (start_ms, end_ms).
_Span = tuple[int, int]


def _check_int(value: object, what: str) -> None:
    # bool is an int subclass; True as a time value is always a bug.
    if not isinstance(value, int) or isinstance(value, bool):
        raise ValueError(f"{what} must be an int number of milliseconds, got {value!r}")


def _check_span(start_ms: int, end_ms: int, what: str) -> None:
    _check_int(start_ms, f"{what} file_start_ms")
    _check_int(end_ms, f"{what} file_end_ms")
    if start_ms < 0 or end_ms <= start_ms:
        raise ValueError(f"{what}: invalid interval {start_ms}..{end_ms} ms")


# --------------------------------------------------------------------------------------
# Public records
# --------------------------------------------------------------------------------------


@dataclass(frozen=True)
class Interval:
    """A non-empty half-open interval ``[file_start_ms, file_end_ms)`` in the source file."""

    file_start_ms: int
    file_end_ms: int

    def __post_init__(self) -> None:
        _check_span(self.file_start_ms, self.file_end_ms, "Interval")

    @property
    def duration_ms(self) -> int:
        return self.file_end_ms - self.file_start_ms


@dataclass(frozen=True)
class SpeakerTurn:
    """A diarization turn: ``speaker`` talks during ``[file_start_ms, file_end_ms)``."""

    speaker: str
    file_start_ms: int
    file_end_ms: int

    def __post_init__(self) -> None:
        if not isinstance(self.speaker, str) or not self.speaker:
            raise ValueError(f"speaker label must be a non-empty string, got {self.speaker!r}")
        _check_span(self.file_start_ms, self.file_end_ms, f"turn of {self.speaker!r}")


@dataclass(frozen=True)
class PlannedSegment:
    """One output segment (padded), cut from the source file.

    Attributes:
        file_start_ms: Padded start, relative to the file.
        file_end_ms: Padded end (exclusive), relative to the file.
        speaker: Diarization label of the only speaker in the segment.
        speech_ms: Milliseconds of speech inside the segment (silences excluded).
    """

    file_start_ms: int
    file_end_ms: int
    speaker: str
    speech_ms: int

    def __post_init__(self) -> None:
        _check_span(self.file_start_ms, self.file_end_ms, "PlannedSegment")
        _check_int(self.speech_ms, "PlannedSegment speech_ms")
        if not self.speaker:
            raise ValueError("PlannedSegment speaker must be non-empty")
        if not 0 < self.speech_ms <= self.file_end_ms - self.file_start_ms:
            raise ValueError(
                f"speech_ms {self.speech_ms} outside 1..{self.file_end_ms - self.file_start_ms}"
            )

    @property
    def duration_ms(self) -> int:
        return self.file_end_ms - self.file_start_ms


@dataclass(frozen=True)
class RejectedSpan:
    """A stretch of speech that is not part of any planned segment, and why.

    Attributes:
        file_start_ms: Start of the rejected speech extent.
        file_end_ms: End (exclusive) of the rejected speech extent.
        reason: One of :data:`REJECT_REASONS`.
        speaker: The speaker, when known (``None`` for ``overlap`` / ``no_speaker``).
    """

    file_start_ms: int
    file_end_ms: int
    reason: RejectReason
    speaker: str | None = None

    def __post_init__(self) -> None:
        _check_span(self.file_start_ms, self.file_end_ms, "RejectedSpan")
        if self.reason not in REJECT_REASONS:
            raise ValueError(f"unknown reject reason {self.reason!r}; expected {REJECT_REASONS}")

    @property
    def duration_ms(self) -> int:
        return self.file_end_ms - self.file_start_ms


@dataclass(frozen=True)
class SegmentPlan:
    """Result of :func:`plan_segments`: both lists are sorted by time."""

    segments: tuple[PlannedSegment, ...]
    rejected: tuple[RejectedSpan, ...]

    def reject_counts(self) -> dict[str, int]:
        """Number of rejected spans per reason (every reason present, possibly 0)."""
        counts = dict.fromkeys(REJECT_REASONS, 0)
        for span in self.rejected:
            counts[span.reason] += 1
        return counts


@dataclass(frozen=True)
class SegmentConfig:
    """S3 thresholds. All values are integer milliseconds; load with :meth:`from_mapping`.

    Attributes:
        min_segment_ms: Minimum speech extent of a segment (before padding).
        max_segment_ms: Maximum length of a segment, padding included.
        min_cut_silence_ms: A cut needs a silence at least this long; shorter pauses are
            merged over.
        pad_ms: Padding per side.
        overlap_guard_ms: Distance kept from overlapped-speech regions.
    """

    min_segment_ms: int
    max_segment_ms: int
    min_cut_silence_ms: int
    pad_ms: int
    overlap_guard_ms: int

    def __post_init__(self) -> None:
        for name in _CONFIG_FIELDS:
            _check_int(getattr(self, name), f"SegmentConfig.{name}")
        if self.min_segment_ms < 1:
            raise ValueError(f"min_segment_ms must be >= 1, got {self.min_segment_ms}")
        if self.max_segment_ms < self.min_segment_ms:
            raise ValueError(
                f"max_segment_ms ({self.max_segment_ms}) must be >= "
                f"min_segment_ms ({self.min_segment_ms})"
            )
        if self.min_cut_silence_ms < 1:
            raise ValueError(f"min_cut_silence_ms must be >= 1, got {self.min_cut_silence_ms}")
        if self.pad_ms < 0:
            raise ValueError(f"pad_ms must be >= 0, got {self.pad_ms}")
        if self.overlap_guard_ms < 0:
            raise ValueError(f"overlap_guard_ms must be >= 0, got {self.overlap_guard_ms}")

    @classmethod
    def from_mapping(cls, mapping: Mapping[str, object]) -> SegmentConfig:
        """Build a config from a parsed YAML mapping.

        Raises:
            ValueError: If ``mapping`` is not a mapping, has unknown or missing keys, or a
                value is not an int (bools and floats are rejected) or violates its bound.
        """
        if not isinstance(mapping, Mapping):
            raise ValueError(f"segment config must be a mapping, got {type(mapping).__name__}")
        keys = set(mapping)
        unknown = sorted(str(key) for key in keys - set(_CONFIG_FIELDS))
        missing = sorted(set(_CONFIG_FIELDS) - keys)
        if unknown or missing:
            raise ValueError(f"segment config: unknown keys {unknown}, missing keys {missing}")
        values: dict[str, int] = {}
        for name in _CONFIG_FIELDS:
            value = mapping[name]
            if not isinstance(value, int) or isinstance(value, bool):
                raise ValueError(f"segment config: {name} must be an int (ms), got {value!r}")
            values[name] = value
        return cls(**values)


_CONFIG_FIELDS: tuple[str, ...] = (
    "min_segment_ms",
    "max_segment_ms",
    "min_cut_silence_ms",
    "pad_ms",
    "overlap_guard_ms",
)


# --------------------------------------------------------------------------------------
# Interval algebra on normalised span lists (sorted, disjoint, non-touching)
# --------------------------------------------------------------------------------------


def _normalize(spans: Iterable[_Span]) -> list[_Span]:
    """Sort and union spans; overlapping and touching spans are merged."""
    out: list[_Span] = []
    for start, end in sorted(spans):
        if out and start <= out[-1][1]:
            if end > out[-1][1]:
                out[-1] = (out[-1][0], end)
        else:
            out.append((start, end))
    return out


def _intersect(a: Sequence[_Span], b: Sequence[_Span]) -> list[_Span]:
    out: list[_Span] = []
    i = j = 0
    while i < len(a) and j < len(b):
        start = max(a[i][0], b[j][0])
        end = min(a[i][1], b[j][1])
        if start < end:
            out.append((start, end))
        if a[i][1] < b[j][1]:
            i += 1
        else:
            j += 1
    return out


def _subtract(a: Sequence[_Span], b: Sequence[_Span]) -> list[_Span]:
    """Parts of ``a`` not covered by ``b``."""
    out: list[_Span] = []
    j = 0
    for start, end in a:
        cur = start
        while j < len(b) and b[j][1] <= cur:
            j += 1
        k = j
        while k < len(b) and b[k][0] < end:
            if b[k][0] > cur:
                out.append((cur, b[k][0]))
            cur = max(cur, b[k][1])
            k += 1
            if cur >= end:
                break
        if cur < end:
            out.append((cur, end))
    return out


def _expand(spans: Iterable[_Span], by_ms: int, file_duration_ms: int) -> list[_Span]:
    return _normalize(
        (max(0, start - by_ms), min(file_duration_ms, end + by_ms)) for start, end in spans
    )


class _SpanIndex:
    """Binary-search queries over a normalised span list."""

    def __init__(self, spans: Sequence[_Span]) -> None:
        self.spans: tuple[_Span, ...] = tuple(spans)
        self._starts = [start for start, _ in spans]
        self._ends = [end for _, end in spans]

    def intersects(self, start_ms: int, end_ms: int) -> bool:
        """True if any span shares at least one millisecond with ``[start_ms, end_ms)``."""
        idx = bisect_right(self._ends, start_ms)  # first span ending after start_ms
        return idx < len(self._starts) and self._starts[idx] < end_ms

    def containing(self, t_ms: int) -> _Span | None:
        """The span with ``start <= t_ms < end``, if any."""
        idx = bisect_right(self._starts, t_ms) - 1
        if idx >= 0 and t_ms < self._ends[idx]:
            return self.spans[idx]
        return None

    def last_end_at_or_before(self, t_ms: int) -> int | None:
        idx = bisect_right(self._ends, t_ms) - 1
        return self._ends[idx] if idx >= 0 else None

    def first_start_at_or_after(self, t_ms: int) -> int | None:
        idx = bisect_left(self._starts, t_ms)
        return self._starts[idx] if idx < len(self._starts) else None


# --------------------------------------------------------------------------------------
# Planning
# --------------------------------------------------------------------------------------


@dataclass(frozen=True)
class _Context:
    """Normalised inputs shared by all planning steps."""

    file_duration_ms: int
    speech: _SpanIndex
    forbidden: _SpanIndex  # overlap regions expanded by the guard
    others: dict[str, _SpanIndex]  # per speaker: union of the *other* speakers' turns


@dataclass(frozen=True)
class _Chunk:
    start_ms: int
    end_ms: int
    speech_ms: int


@dataclass(frozen=True)
class _Candidate:
    speaker: str
    start_ms: int
    end_ms: int
    speech_ms: int


def _spans_of(intervals: Sequence[Interval], what: str, file_duration_ms: int) -> list[_Span]:
    spans: list[_Span] = []
    for item in intervals:
        if not isinstance(item, Interval):
            raise TypeError(f"{what} must contain Interval objects, got {type(item).__name__}")
        if item.file_end_ms > file_duration_ms:
            raise ValueError(
                f"{what}: interval {item.file_start_ms}..{item.file_end_ms} ms exceeds the "
                f"file duration {file_duration_ms} ms"
            )
        spans.append((item.file_start_ms, item.file_end_ms))
    return _normalize(spans)


def _turns_by_speaker(
    turns: Sequence[SpeakerTurn], file_duration_ms: int
) -> dict[str, list[_Span]]:
    raw: dict[str, list[_Span]] = {}
    for turn in turns:
        if not isinstance(turn, SpeakerTurn):
            raise TypeError(f"turns must contain SpeakerTurn objects, got {type(turn).__name__}")
        if turn.file_end_ms > file_duration_ms:
            raise ValueError(
                f"turn of {turn.speaker!r}: {turn.file_start_ms}..{turn.file_end_ms} ms "
                f"exceeds the file duration {file_duration_ms} ms"
            )
        raw.setdefault(turn.speaker, []).append((turn.file_start_ms, turn.file_end_ms))
    return {speaker: _normalize(spans) for speaker, spans in sorted(raw.items())}


def _classify_speech(
    ctx: _Context, speech: list[_Span], turns: dict[str, list[_Span]]
) -> tuple[dict[str, list[_Span]], list[RejectedSpan]]:
    """Step (a): split speech into per-speaker regions and rejected pieces.

    Every millisecond of ``speech`` lands in exactly one place: a speaker region, an
    ``overlap`` rejection (inside the forbidden zone, whoever speaks) or a ``no_speaker``
    rejection (outside every turn).
    """
    rejected: list[RejectedSpan] = []
    forbidden = list(ctx.forbidden.spans)
    for start, end in _intersect(speech, forbidden):
        rejected.append(RejectedSpan(start, end, "overlap"))
    rest = _subtract(speech, forbidden)
    any_turn = _normalize(span for spans in turns.values() for span in spans)
    for start, end in _subtract(rest, any_turn):
        rejected.append(RejectedSpan(start, end, "no_speaker"))
    regions = {speaker: _intersect(rest, spans) for speaker, spans in turns.items()}
    return regions, rejected


def _split_runs(regions: Sequence[_Span], speaker: str, ctx: _Context) -> list[list[_Span]]:
    """Split one speaker's regions into runs at *walls*.

    Between two consecutive regions the gap is a wall if anything but plain silence lies in
    it: other speech (e.g. a rejected piece), the forbidden zone, or another speaker's
    turn. Segments never span a wall, so a segment never contains foreign material.
    """
    runs: list[list[_Span]] = []
    others = ctx.others[speaker]
    for region in regions:
        if runs:
            gap_start = runs[-1][-1][1]
            gap_end = region[0]
            wall = (
                ctx.speech.intersects(gap_start, gap_end)
                or ctx.forbidden.intersects(gap_start, gap_end)
                or others.intersects(gap_start, gap_end)
            )
            if not wall:
                runs[-1].append(region)
                continue
        runs.append([region])
    return runs


def _merge_chunks(run: Sequence[_Span], cfg: SegmentConfig) -> list[_Chunk]:
    """Step (b): merge regions separated by silences shorter than ``min_cut_silence_ms``."""
    chunks: list[_Chunk] = []
    for start, end in run:
        if chunks and start - chunks[-1].end_ms < cfg.min_cut_silence_ms:
            last = chunks[-1]
            chunks[-1] = _Chunk(last.start_ms, end, last.speech_ms + (end - start))
        else:
            chunks.append(_Chunk(start, end, end - start))
    return chunks


def _pack_run(
    chunks: Sequence[_Chunk], speaker: str, cfg: SegmentConfig
) -> tuple[list[_Candidate], list[RejectedSpan]]:
    """Steps (c) and (e): greedy packing, then the minimum-length check."""
    candidates: list[_Candidate] = []
    rejected: list[RejectedSpan] = []
    i = 0
    while i < len(chunks):
        first = chunks[i]
        if first.end_ms - first.start_ms > cfg.max_segment_ms:
            rejected.append(
                RejectedSpan(first.start_ms, first.end_ms, "too_long_without_pause", speaker)
            )
            i += 1
            continue
        j = i + 1
        speech_ms = first.speech_ms
        while j < len(chunks) and chunks[j].end_ms - first.start_ms <= cfg.max_segment_ms:
            speech_ms += chunks[j].speech_ms
            j += 1
        end_ms = chunks[j - 1].end_ms
        if end_ms - first.start_ms < cfg.min_segment_ms:
            rejected.append(RejectedSpan(first.start_ms, end_ms, "too_short", speaker))
        else:
            candidates.append(_Candidate(speaker, first.start_ms, end_ms, speech_ms))
        i = j
    return candidates, rejected


def _fit_padding(pad_left_ms: int, pad_right_ms: int, room_ms: int) -> tuple[int, int]:
    """Shrink the two paddings so that their sum is at most ``room_ms`` (fairly)."""
    if pad_left_ms + pad_right_ms <= room_ms:
        return pad_left_ms, pad_right_ms
    half = room_ms // 2
    if pad_left_ms <= half:
        return pad_left_ms, room_ms - pad_left_ms
    if pad_right_ms <= room_ms - half:
        return room_ms - pad_right_ms, pad_right_ms
    return half, room_ms - half


def _left_limit_ms(start_ms: int, speaker: str, ctx: _Context) -> int:
    """Smallest start the padded segment may take (a segment starts inside speech)."""
    containing = ctx.speech.containing(start_ms)
    if containing is not None and containing[0] < start_ms:
        return start_ms  # speech continues to the left: the cut is not in a silence
    limit = 0
    previous_end = ctx.speech.last_end_at_or_before(start_ms)
    if previous_end is not None:
        limit = max(limit, previous_end + (start_ms - previous_end) // 2)
    for index in (ctx.forbidden, ctx.others[speaker]):
        wall = index.last_end_at_or_before(start_ms)
        if wall is not None:
            limit = max(limit, wall)
    return min(limit, start_ms)


def _right_limit_ms(end_ms: int, speaker: str, ctx: _Context) -> int:
    """Largest end the padded segment may take."""
    containing = ctx.speech.containing(end_ms - 1)
    if containing is not None and end_ms < containing[1]:
        return end_ms  # speech continues to the right
    limit = ctx.file_duration_ms
    next_start = ctx.speech.first_start_at_or_after(end_ms)
    if next_start is not None:
        limit = min(limit, end_ms + (next_start - end_ms) // 2)
    for index in (ctx.forbidden, ctx.others[speaker]):
        wall = index.first_start_at_or_after(end_ms)
        if wall is not None:
            limit = min(limit, wall)
    return max(limit, end_ms)


def _pad(candidate: _Candidate, ctx: _Context, cfg: SegmentConfig) -> PlannedSegment:
    """Step (d): pad both sides, then trim so the segment fits ``max_segment_ms``."""
    pad_left = min(
        cfg.pad_ms, candidate.start_ms - _left_limit_ms(candidate.start_ms, candidate.speaker, ctx)
    )
    pad_right = min(
        cfg.pad_ms, _right_limit_ms(candidate.end_ms, candidate.speaker, ctx) - candidate.end_ms
    )
    room_ms = cfg.max_segment_ms - (candidate.end_ms - candidate.start_ms)
    pad_left, pad_right = _fit_padding(pad_left, pad_right, room_ms)
    return PlannedSegment(
        file_start_ms=candidate.start_ms - pad_left,
        file_end_ms=candidate.end_ms + pad_right,
        speaker=candidate.speaker,
        speech_ms=candidate.speech_ms,
    )


def plan_segments(
    file_duration_ms: int,
    speech: Sequence[Interval],
    turns: Sequence[SpeakerTurn],
    overlaps: Sequence[Interval],
    cfg: SegmentConfig,
) -> SegmentPlan:
    """Plan the segments of one source file (S3).

    All times are integer milliseconds relative to the start of the file. Inputs may be
    unsorted; overlapping or touching intervals of the same list are unioned. Wherever two
    different speakers' turns intersect, that region counts as overlapped speech too, even
    if ``overlaps`` does not list it. Speech that lies outside every turn is rejected with
    ``no_speaker``; speech inside an overlap region (expanded by ``cfg.overlap_guard_ms``)
    with ``overlap``.

    Guarantees (property-tested):
        * segments are inside ``[0, file_duration_ms]``, pairwise disjoint (they may touch),
          single-speaker, and ``cfg.min_segment_ms <= speech extent`` and
          ``duration_ms <= cfg.max_segment_ms``;
        * no segment intersects the expanded overlap regions or another speaker's turn;
        * every millisecond of speech lies in exactly one segment or one rejected span;
        * the result is deterministic and independent of input order.

    Note:
        A cut forced by a speaker change or by the overlap guard can fall inside speech
        (there is no silence to cut in); the padding on that side is then zero. Half-words
        at such edges are expected to be caught by the S4 / S5 filters.

    Args:
        file_duration_ms: Length of the source file; must be > 0.
        speech: VAD speech intervals.
        turns: Diarization turns (speaker label + interval).
        overlaps: Regions where two or more speakers talk.
        cfg: Segmentation thresholds.

    Returns:
        A :class:`SegmentPlan` with segments and rejected spans sorted by time.

    Raises:
        TypeError: If a list element is not an :class:`Interval` / :class:`SpeakerTurn`.
        ValueError: If ``file_duration_ms`` is not a positive int, or any interval lies outside
            ``[0, file_duration_ms]`` (negative and empty intervals are rejected when the
            :class:`Interval` / :class:`SpeakerTurn` is constructed).
    """
    _check_int(file_duration_ms, "file_duration_ms")
    if file_duration_ms <= 0:
        raise ValueError(f"file_duration_ms must be > 0, got {file_duration_ms}")
    speech_spans = _spans_of(speech, "speech", file_duration_ms)
    turn_spans = _turns_by_speaker(turns, file_duration_ms)
    overlap_spans = _spans_of(overlaps, "overlaps", file_duration_ms)

    speakers = list(turn_spans)
    for index, first in enumerate(speakers):
        for second in speakers[index + 1 :]:
            overlap_spans.extend(_intersect(turn_spans[first], turn_spans[second]))
    forbidden = _expand(_normalize(overlap_spans), cfg.overlap_guard_ms, file_duration_ms)

    ctx = _Context(
        file_duration_ms=file_duration_ms,
        speech=_SpanIndex(speech_spans),
        forbidden=_SpanIndex(forbidden),
        others={
            speaker: _SpanIndex(
                _normalize(
                    span for other in speakers if other != speaker for span in turn_spans[other]
                )
            )
            for speaker in speakers
        },
    )

    regions, rejected = _classify_speech(ctx, speech_spans, turn_spans)
    candidates: list[_Candidate] = []
    for speaker, speaker_regions in regions.items():
        for run in _split_runs(speaker_regions, speaker, ctx):
            run_candidates, run_rejected = _pack_run(_merge_chunks(run, cfg), speaker, cfg)
            candidates.extend(run_candidates)
            rejected.extend(run_rejected)

    segments = sorted(
        (_pad(candidate, ctx, cfg) for candidate in candidates),
        key=lambda seg: (seg.file_start_ms, seg.file_end_ms, seg.speaker),
    )
    rejected.sort(key=lambda span: (span.file_start_ms, span.file_end_ms, span.reason))
    return SegmentPlan(segments=tuple(segments), rejected=tuple(rejected))