backend/alignment.py

"""Compare scripted dialogue to the processed waveform (channels mode) and emit
timestamped error records for missing / misaligned / extra speech.

Method (VAD, no ASR):
  * Each character has an isolated channel WAV, so VAD tells us *when that voice
    actually speaks*.
  * For every scripted line we check whether speech is present in the right place:
      - no speech where the script expects it          -> MISSING
      - speech present but its start/end drifts > tol   -> MISALIGNED (onset/offset)
      - speech much shorter than the scripted span      -> MISALIGNED/truncated
      - VAD speech with no scripted line for that char   -> EXTRA
  * A constant capture offset (script-TC zero vs audio zero) would otherwise mark
    every line misaligned, so we ESTIMATE the per-channel offset and subtract it
    before scoring drift.
"""

from __future__ import annotations

from pathlib import Path
from statistics import median
from typing import Any

from pydantic import BaseModel

from .characters import CharacterEntity
from .script_parser import ScriptDoc
from .vad import detect_speech_regions

Interval = tuple[float, float]


class AlignError(BaseModel):
    type: str                       # MISSING | MISALIGNED | EXTRA
    subtype: str | None = None      # onset_drift | offset_drift | truncated
    severity: str                   # error | warn | info
    character: str | None = None
    channel: str | None = None
    script_index: int | None = None
    script_start_s: float | None = None
    script_end_s: float | None = None
    audio_start_s: float | None = None
    audio_end_s: float | None = None
    drift_s: float | None = None
    coverage: float | None = None   # fraction of the scripted span covered by speech
    text: str | None = None         # the scripted dialogue line (for MISSING/MISALIGNED)
    message: str = ""


class ChannelAlignment(BaseModel):
    character: str
    channel: str
    offset_s: float                 # estimated script->audio offset applied
    n_lines: int
    n_missing: int
    n_misaligned: int
    n_extra: int
    errors: list[AlignError]


# --------------------------------------------------------------------------- #
# interval helpers
# --------------------------------------------------------------------------- #
def _overlap(a: Interval, b: Interval) -> float:
    return max(0.0, min(a[1], b[1]) - max(a[0], b[0]))


def _coverage(span: Interval, regions: list[Interval]) -> tuple[float, Interval | None]:
    """Return (covered_fraction, merged_matched_span) for `span` against regions."""
    covered = 0.0
    lo = hi = None
    for r in regions:
        ov = _overlap(span, r)
        if ov > 0:
            covered += ov
            lo = r[0] if lo is None else min(lo, r[0])
            hi = r[1] if hi is None else max(hi, r[1])
    dur = max(1e-9, span[1] - span[0])
    return covered / dur, (None if lo is None else (lo, hi))


def estimate_offset(
    script_spans: list[Interval], regions: list[Interval], max_offset_s: float = 5.0
) -> float:
    """Median (nearest-region-onset - script-onset) over lines that have a nearby
    region — robust to missing lines. 0.0 when there's nothing to anchor on."""
    deltas: list[float] = []
    starts = sorted(r[0] for r in regions)
    for s0, _ in script_spans:
        # nearest region start
        best = None
        for rs in starts:
            d = rs - s0
            if abs(d) <= max_offset_s and (best is None or abs(d) < abs(best)):
                best = d
        if best is not None:
            deltas.append(best)
    if len(deltas) < 3:
        return 0.0
    return round(median(deltas), 3)


# --------------------------------------------------------------------------- #
# core
# --------------------------------------------------------------------------- #
def align_channel(
    character: str,
    channel: str,
    script_spans: list[tuple[int, float, float]],   # (script_index, start, end)
    regions: list[Interval],
    *,
    tol_s: float = 0.5,
    missing_coverage: float = 0.15,
    truncated_ratio: float = 0.6,
    offset_s: float | None = None,
) -> ChannelAlignment:
    """Score one character's scripted lines against their channel's VAD regions."""
    spans = [(a, b) for _, a, b in script_spans]
    if offset_s is None:
        offset_s = estimate_offset(spans, regions)

    errors: list[AlignError] = []
    matched_regions: list[Interval] = []
    n_missing = n_misaligned = 0

    for idx, a, b in script_spans:
        span = (a + offset_s, b + offset_s)
        cov, matched = _coverage(span, regions)

        if cov < missing_coverage or matched is None:
            n_missing += 1
            errors.append(AlignError(
                type="MISSING", severity="error", character=character, channel=channel,
                script_index=idx, script_start_s=round(a, 3), script_end_s=round(b, 3),
                coverage=round(cov, 3),
                message=f"No speech in '{channel}' for scripted line {idx} "
                        f"({a:.2f}-{b:.2f}s); coverage {cov:.0%}.",
            ))
            continue

        matched_regions.append(matched)
        onset_drift = matched[0] - span[0]
        offset_drift = matched[1] - span[1]
        span_dur = max(1e-9, span[1] - span[0])
        matched_dur = matched[1] - matched[0]

        subtype = None
        if matched_dur < truncated_ratio * span_dur:
            subtype = "truncated"
        elif abs(onset_drift) > abs(offset_drift) and abs(onset_drift) > tol_s:
            subtype = "onset_drift"
        elif abs(offset_drift) > tol_s:
            subtype = "offset_drift"

        if subtype:
            n_misaligned += 1
            drift = onset_drift if subtype == "onset_drift" else offset_drift
            if subtype == "truncated":
                drift = matched_dur - span_dur
            errors.append(AlignError(
                type="MISALIGNED", subtype=subtype, severity="warn",
                character=character, channel=channel, script_index=idx,
                script_start_s=round(a, 3), script_end_s=round(b, 3),
                audio_start_s=round(matched[0] - offset_s, 3),
                audio_end_s=round(matched[1] - offset_s, 3),
                drift_s=round(drift, 3), coverage=round(cov, 3),
                message=f"Line {idx} {subtype.replace('_', ' ')} by {drift:+.2f}s "
                        f"in '{channel}'.",
            ))

    # EXTRA: speech regions not overlapping any scripted (offset-shifted) line.
    shifted = [(a + offset_s, b + offset_s) for a, b in spans]
    for r in regions:
        if all(_overlap(r, s) <= 0 for s in shifted):
            errors.append(AlignError(
                type="EXTRA", severity="info", character=character, channel=channel,
                audio_start_s=round(r[0] - offset_s, 3), audio_end_s=round(r[1] - offset_s, 3),
                message=f"Speech in '{channel}' at {r[0]:.2f}-{r[1]:.2f}s with no scripted line.",
            ))
    n_extra = sum(1 for e in errors if e.type == "EXTRA")

    return ChannelAlignment(
        character=character, channel=channel, offset_s=offset_s,
        n_lines=len(script_spans), n_missing=n_missing,
        n_misaligned=n_misaligned, n_extra=n_extra, errors=errors,
    )


def align_script_to_channels(
    doc: ScriptDoc,
    characters: list[CharacterEntity],
    channel_wavs: dict[str, Path],       # channel_name -> wav path
    *,
    tol_s: float = 0.5,
    vad_kwargs: dict[str, Any] | None = None,
    offset_s: float | None = None,
) -> dict[str, Any]:
    """Full pass: per character with a mapped channel, VAD the channel and score
    their lines. Returns a JSON-serialisable report."""
    spans_by_char: dict[str, list[tuple[int, float, float]]] = {}
    for seg in doc.segments:
        for key in seg.characters:
            spans_by_char.setdefault(key, []).append((seg.index, seg.start_s, seg.end_s))

    region_cache: dict[str, list[Interval]] = {}
    channel_reports: list[ChannelAlignment] = []
    unmapped: list[str] = []

    for ent in characters:
        spans = spans_by_char.get(ent.id, [])
        if not spans:
            continue
        if not ent.channel or ent.channel not in channel_wavs:
            unmapped.append(ent.id)
            continue
        if ent.channel not in region_cache:
            regs = detect_speech_regions(channel_wavs[ent.channel], **(vad_kwargs or {}))
            region_cache[ent.channel] = [(r["start"], r["end"]) for r in regs]
        channel_reports.append(align_channel(
            ent.id, ent.channel, spans, region_cache[ent.channel],
            tol_s=tol_s, offset_s=offset_s,
        ))

    # Attach the scripted dialogue line to each error that references a script index.
    text_by_index = {seg.index: seg.text for seg in doc.segments}
    for cr in channel_reports:
        for e in cr.errors:
            if e.script_index is not None:
                e.text = text_by_index.get(e.script_index)

    all_errors = [e for cr in channel_reports for e in cr.errors]
    return {
        "tol_s": tol_s,
        "channels": [cr.model_dump() for cr in channel_reports],
        "errors": [e.model_dump() for e in all_errors],
        "unmapped_characters": unmapped,
        "summary": {
            "n_characters_checked": len(channel_reports),
            "n_missing": sum(cr.n_missing for cr in channel_reports),
            "n_misaligned": sum(cr.n_misaligned for cr in channel_reports),
            "n_extra": sum(cr.n_extra for cr in channel_reports),
            "n_unmapped": len(unmapped),
        },
    }