"""Speaker-similarity measurements on top of the ECAPA-TDNN embedding worker.

Two uses, driven by ``scripts/speaker_sim.py``:

* **say**: cosine similarity of every generated chunk of a ``nida say`` output directory to
  the embedding of its speaker prompt (:func:`read_say_chunks`, :func:`resolve_references`,
  :func:`score_say`, :func:`summarize_say`);
* **groups**: does a dataset ``group`` (speaker / source video) really hold one voice? Intra-
  group versus inter-group cosine similarity of sampled segments (:func:`read_index`,
  :func:`select_candidates`, :func:`sample_groups`, :func:`summarize_groups`).

Embedding is a side effect at the edge (:func:`embed_with_adapter`: the worker protocol with
its resumable results file, run in-process); everything else is pure and needs only numpy.
Embeddings are 192-d float32 unit vectors, so a cosine similarity is a dot product. Thresholds
are arguments (the script's flags), never constants of the measurement.

``say`` layout relied upon (``nida.infer.say``): ``say.jsonl`` (``prompt_id``, ``index`` per
chunk), ``chunks/<prompt_id>__<index:02d>.16k.wav`` (the 16 kHz generated chunk),
``settings.json`` (``external_prompts``: external prompt id -> source wav path as given to
``--prompt-list``, with ``external_prompt_texts`` and ``external_prompt_sha256`` beside it in newer
runs; all absent or empty for dataset prompts) and ``audio/<chunk>.16k.wav`` (the decoder's 16 kHz
output for ``prompt + sentence``: ``say`` cuts the generated chunk from it with
``prompt_frames`` / ``num_frames`` of ``say.jsonl``, and the same cut's first part is the
decoded prompt).
"""

from __future__ import annotations

import hashlib
import io
import json
import logging
import math
import os
import random
from collections import defaultdict
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from math import gcd
from pathlib import Path
from typing import Literal

import numpy as np
import numpy.typing as npt

from nida.bench.protocol import (
    WorkRequest,
    drop_results,
    load_requests,
    load_results,
    read_summary,
    write_requests,
)
from nida.bench.workers.runner import Adapter, ModelRef, WorkerOptions, run_worker
from nida.data.workers.spk_ecapa import EMBEDDING_DIM, sha256_file
from nida.text.turkish import nfc
from nida.utils.io import atomic_write_bytes

__all__ = [
    "DEFAULT_DRIFT_THRESHOLD",
    "ChunkRef",
    "ClipEmbedding",
    "ClipFailure",
    "EmbedOutcome",
    "Embedder",
    "ExternalPromptInfo",
    "IndexRow",
    "PromptListEntry",
    "ReferenceMode",
    "ReferenceRef",
    "SpeakerSimInputError",
    "cosine",
    "cut_prompt_audio",
    "describe",
    "embed_with_adapter",
    "find_dataset_dir",
    "list_audio_stems",
    "prompt_id_for",
    "read_external_prompts",
    "read_index",
    "read_prompt_list",
    "read_say_chunks",
    "read_say_context",
    "resolve_references",
    "sample_groups",
    "score_say",
    "select_candidates",
    "summarize_groups",
    "summarize_say",
    "unmatched_prompt_entries",
    "validate_group_options",
    "validate_threshold",
    "write_16k_mono_wav",
    "write_decoded_prompt_wav",
    "write_if_changed",
]

log = logging.getLogger(__name__)

#: Default for ``--threshold`` of the ``say`` subcommand (the user's choice, 2026-10-15).
DEFAULT_DRIFT_THRESHOLD = 0.5
#: Sample rate the embedding worker requires.
WORK_RATE_HZ = 16000
#: Unit-norm tolerance for embeddings read back from disk.
NORM_TOLERANCE = 1e-3
SAY_MANIFEST = "say.jsonl"
SAY_SETTINGS = "settings.json"
LOWEST_CHUNKS = 5

ReferenceMode = Literal["auto", "source", "decoded"]
FloatVec = npt.NDArray[np.float32]


class SpeakerSimInputError(ValueError):
    """An input directory, manifest or option is missing, malformed or inconsistent."""


# ------------------------------------------------------------------------------ embedding


@dataclass(frozen=True)
class ClipEmbedding:
    """A unit-norm embedding ``(192,)`` and the duration of the clip it came from."""

    vector: FloatVec
    duration_ms: float


@dataclass(frozen=True)
class ClipFailure:
    """A clip the worker could not embed (``"<ExceptionType>: <message>"``)."""

    error: str


EmbedOutcome = ClipEmbedding | ClipFailure
#: Embeds the clips of the requests, keeping its files (results, embeddings) in the given work
#: directory (:func:`embed_with_adapter`): ``clip_id -> outcome``.
Embedder = Callable[[Sequence[WorkRequest], Path], Mapping[str, EmbedOutcome]]


def embed_with_adapter(
    adapter: Adapter,
    model: ModelRef,
    requests: Sequence[WorkRequest],
    work_dir: Path,
    *,
    device: str,
    batch_size: int,
) -> dict[str, EmbedOutcome]:
    """Run ``adapter`` over ``requests`` through the worker protocol and read the embeddings.

    ``work_dir`` holds ``requests.jsonl``, ``results.jsonl``, ``results.summary.json`` and
    ``embeddings/<clip_id>.npy``. The run is resumable: results of an earlier invocation are
    reused when their embedding file exists and the clip's audio is unchanged (each request
    carries an ``audio_fingerprint`` param, size and mtime of the file, so a ``say`` rerun into
    the same directory is never served stale embeddings, and a ``model`` param ``id@revision``);
    failed clips are tried again, results of clips no longer requested, changed or lost are
    dropped. Results of another model / revision are refused. Each request's ``out_path`` is set
    here.

    Raises:
        SpeakerSimInputError: results of another model, no result for a request, or an
            embedding file that is not a unit vector of the expected shape.
        TooManyFailuresError: (from the worker loop) the model or device is broken.
    """
    work_dir.mkdir(parents=True, exist_ok=True)
    emb_dir = work_dir / "embeddings"
    prepared = [
        WorkRequest(
            r.clip_id,
            r.audio_path,
            {
                **r.params,
                "out_path": str(emb_dir / f"{r.clip_id}.npy"),
                "audio_fingerprint": _fingerprint(r.audio_path),
                "model": f"{model.model_id}@{model.revision}",
            },
        )
        for r in requests
    ]
    requests_path = work_dir / "requests.jsonl"
    results_path = work_dir / "results.jsonl"
    summary_path = work_dir / "results.summary.json"
    if results_path.exists() and summary_path.exists():
        previous = read_summary(summary_path)
        if (previous.model_id, previous.revision) != (model.model_id, model.revision):
            raise SpeakerSimInputError(
                f"{work_dir}: existing results come from {previous.model_id}@{previous.revision}, "
                f"not {model.model_id}@{model.revision}; move them away first"
            )
    if results_path.exists():
        # Without the previous requests nothing proves a result belongs to the current audio.
        before = (
            {r.clip_id: r for r in load_requests(requests_path)} if requests_path.exists() else {}
        )
        current = {r.clip_id: r for r in prepared}
        existing = load_results(results_path, repair=True)
        doomed = [
            cid
            for cid, res in existing.items()
            if cid not in current
            or before.get(cid) != current[cid]
            or res.output is None
            or not Path(str(res.output.get("out_path", ""))).is_file()
        ]
        if doomed:
            log.info("dropping %d stray / failed / lost results before the run", len(doomed))
            drop_results(results_path, doomed)
    write_requests(requests_path, prepared)
    run_worker(
        adapter,
        WorkerOptions(
            requests=requests_path,
            results=results_path,
            summary=summary_path,
            model=model,
            device=device,
            batch_size=batch_size,
            retry_failed=False,
        ),
    )
    results = load_results(results_path)
    outcomes: dict[str, EmbedOutcome] = {}
    for request in prepared:
        result = results.get(request.clip_id)
        if result is None:
            raise SpeakerSimInputError(f"{request.clip_id}: the worker produced no result")
        if result.output is None:
            outcomes[request.clip_id] = ClipFailure(str(result.error))
            continue
        outcomes[request.clip_id] = _read_embedding(request.clip_id, result.output)
    return outcomes


def _fingerprint(audio_path: str) -> str:
    """``<size>:<mtime_ns>`` of an audio file: changes whenever the file is rewritten."""
    try:
        stat = Path(audio_path).stat()
    except OSError as exc:
        raise SpeakerSimInputError(f"{audio_path}: cannot stat the audio ({exc})") from exc
    return f"{stat.st_size}:{stat.st_mtime_ns}"


def _read_embedding(clip_id: str, output: Mapping[str, object]) -> ClipEmbedding:
    path = Path(str(output["out_path"]))
    vector = np.load(path, allow_pickle=False)
    if vector.shape != (EMBEDDING_DIM,) or vector.dtype != np.float32:
        raise SpeakerSimInputError(
            f"{clip_id}: {path} has shape {vector.shape} / dtype {vector.dtype}, "
            f"expected ({EMBEDDING_DIM},) float32"
        )
    norm = float(np.linalg.norm(vector))
    if not math.isfinite(norm) or abs(norm - 1.0) > NORM_TOLERANCE:
        raise SpeakerSimInputError(f"{clip_id}: {path} is not a unit vector (norm {norm})")
    duration = output.get("duration_ms")
    if isinstance(duration, bool) or not isinstance(duration, int | float):
        raise SpeakerSimInputError(f"{clip_id}: worker output lacks duration_ms")
    return ClipEmbedding(vector, float(duration))


# --------------------------------------------------------------------------- numeric helpers


def cosine(a: npt.ArrayLike, b: npt.ArrayLike) -> float:
    """Cosine similarity of two non-zero finite vectors of equal shape.

    Raises:
        ValueError: on a shape mismatch, non-finite values or a zero vector.
    """
    x, y = np.asarray(a, dtype=np.float64), np.asarray(b, dtype=np.float64)
    if x.shape != y.shape or x.ndim != 1:
        raise ValueError(f"cosine needs two 1-D vectors of one length, got {x.shape} and {y.shape}")
    if not (np.isfinite(x).all() and np.isfinite(y).all()):
        raise ValueError("cosine: non-finite values")
    nx, ny = float(np.linalg.norm(x)), float(np.linalg.norm(y))
    if nx == 0.0 or ny == 0.0:
        raise ValueError("cosine: zero vector")
    return float(np.dot(x, y) / (nx * ny))


def describe(values: Sequence[float]) -> dict[str, float | int | None]:
    """``n``, ``mean``, ``std`` (sample, ``ddof=1``; 0.0 for one value), ``min``, ``max`` and the
    5/25/50/75/95th percentiles (linear interpolation) of ``values``; ``None`` fields if empty.

    Raises:
        ValueError: on non-finite values.
    """
    keys = ("mean", "std", "min", "max", "p05", "p25", "p50", "p75", "p95")
    if not values:
        return {"n": 0, **dict.fromkeys(keys)}
    data = np.asarray(values, dtype=np.float64)
    if not np.isfinite(data).all():
        raise ValueError("describe: non-finite values")
    p05, p25, p50, p75, p95 = (float(x) for x in np.percentile(data, [5, 25, 50, 75, 95]))
    return {
        "n": int(data.size),
        "mean": float(data.mean()),
        "std": float(data.std(ddof=1)) if data.size > 1 else 0.0,
        "min": float(np.min(data)),
        "max": float(np.max(data)),
        "p05": p05,
        "p25": p25,
        "p50": p50,
        "p75": p75,
        "p95": p95,
    }


def write_if_changed(path: Path, data: bytes) -> bool:
    """Atomically write ``data`` to ``path`` unless it already holds exactly these bytes.

    Returns ``True`` if the file was written. An identical file is left alone so that its
    mtime stays put: the embedding loop fingerprints audio by size and mtime, and a rewrite
    would make a rerun redo work that is still valid.
    """
    if path.is_file() and path.read_bytes() == data:
        return False
    atomic_write_bytes(path, data)
    return True


def _wav_bytes(audio: FloatVec) -> bytes:
    import soundfile  # type: ignore[import-untyped]  # bench extra

    buffer = io.BytesIO()
    soundfile.write(buffer, np.clip(audio, -1.0, 1.0), WORK_RATE_HZ, format="WAV", subtype="PCM_16")
    return buffer.getvalue()


def write_16k_mono_wav(source: Path, target: Path) -> float:
    """Write a mono 16 kHz 16-bit PCM copy of ``source`` (unless identical); returns its
    duration (ms).

    Channels are averaged and the signal resampled with a polyphase filter, exactly as
    ``nida.infer.external_prompt`` does for its prompts. Needs soundfile and scipy.

    Raises:
        SpeakerSimInputError: if the file is empty.
    """
    import soundfile

    audio, rate = soundfile.read(str(source), dtype="float32", always_2d=True)
    mono = audio.mean(axis=1).astype(np.float32)
    if mono.size == 0:
        raise SpeakerSimInputError(f"{source}: empty audio")
    if int(rate) != WORK_RATE_HZ:
        from scipy import signal  # type: ignore[import-untyped]  # bench extra

        g = gcd(int(rate), WORK_RATE_HZ)
        mono = signal.resample_poly(mono, WORK_RATE_HZ // g, int(rate) // g).astype(np.float32)
    write_if_changed(target, _wav_bytes(mono))
    return float(1000.0 * mono.shape[0] / WORK_RATE_HZ)


def cut_prompt_audio(audio: FloatVec, prompt_frames: int, num_frames: int) -> FloatVec:
    """The prompt part of a decoded ``prompt + sentence`` item: the first
    ``round(prompt_frames * len(audio) / num_frames)`` samples, the cut of ``nida.infer.say._cut``
    (a test pins equality) that also produced the generated chunks from the same file.

    Raises:
        SpeakerSimInputError: unless ``0 < prompt_frames <= num_frames``.
    """
    if not 0 < prompt_frames <= num_frames:
        raise SpeakerSimInputError(
            f"need 0 < prompt_frames <= num_frames, got {prompt_frames} / {num_frames}"
        )
    at = round(prompt_frames * audio.shape[0] / num_frames)
    return audio[:at]


def write_decoded_prompt_wav(
    audio_path: Path, prompt_frames: int, num_frames: int, target: Path
) -> float:
    """Write the prompt part of a decoded ``say`` item (``audio/<chunk>.16k.wav``, 16 kHz) as
    the reference wav (unless identical); returns its duration (ms).

    Raises:
        SpeakerSimInputError: not 16 kHz audio, bad frame counts, or an empty prompt part.
    """
    import soundfile

    audio, rate = soundfile.read(str(audio_path), dtype="float32")
    if int(rate) != WORK_RATE_HZ:
        raise SpeakerSimInputError(f"{audio_path}: expected {WORK_RATE_HZ} Hz, got {rate} Hz")
    if audio.ndim != 1:
        audio = audio.mean(axis=1).astype(np.float32)
    prompt = cut_prompt_audio(audio, prompt_frames, num_frames)
    if prompt.size == 0:
        raise SpeakerSimInputError(f"{audio_path}: the prompt part is empty")
    write_if_changed(target, _wav_bytes(prompt))
    return float(1000.0 * prompt.shape[0] / WORK_RATE_HZ)


def validate_threshold(name: str, value: float) -> None:
    """Raise unless ``value`` is a finite number (a cosine threshold).

    Raises:
        SpeakerSimInputError: on NaN or infinity.
    """
    if not math.isfinite(value):
        raise SpeakerSimInputError(f"{name} must be a finite number, got {value}")


def validate_group_options(
    *,
    n_groups: int,
    per_group: int,
    inter_pairs: int,
    min_duration_ms: float,
    intra_threshold: float,
) -> None:
    """Check the ``groups`` options before any model work.

    Raises:
        SpeakerSimInputError: ``n_groups < 1``, ``per_group < 2`` (one segment has no intra
            pair), ``inter_pairs < 1``, a negative / non-finite ``min_duration_ms`` or a
            non-finite ``intra_threshold``.
    """
    if n_groups < 1:
        raise SpeakerSimInputError(f"--groups must be >= 1, got {n_groups}")
    if per_group < 2:
        raise SpeakerSimInputError(f"--per-group must be >= 2, got {per_group}")
    if inter_pairs < 1:
        raise SpeakerSimInputError(f"--inter-pairs must be >= 1, got {inter_pairs}")
    if not math.isfinite(min_duration_ms) or min_duration_ms < 0:
        raise SpeakerSimInputError(f"--min-duration-ms must be >= 0, got {min_duration_ms}")
    validate_threshold("--intra-threshold", intra_threshold)


# ------------------------------------------------------------------------------- say: input


@dataclass(frozen=True)
class ChunkRef:
    """One generated chunk of a ``say`` run.

    ``prompt_frames`` / ``num_frames`` are the item's frame counts from ``say.jsonl`` (``None``
    in a manifest that does not have them); they locate the prompt inside the decoded audio.
    """

    chunk_id: str
    prompt_id: str
    index: int
    wav_path: Path
    prompt_frames: int | None = None
    num_frames: int | None = None


@dataclass(frozen=True)
class ReferenceRef:
    """The audio a prompt's chunks are compared with.

    Attributes:
        kind: ``"source"`` (the recording the prompt came from) or ``"decoded"`` (the prompt
            part of the decoded audio the model's output was cut from).
        wav_path: The file: a source recording (any rate / channel count), or for ``decoded``
            the 16 kHz ``audio/<chunk>.16k.wav`` of the prompt's first chunk.
        prompt_frames: ``decoded`` only: latent frames of the prompt ...
        num_frames: ... and of the whole item; the prompt is the first
            ``prompt_frames / num_frames`` of the file (:func:`cut_prompt_audio`).
    """

    kind: Literal["source", "decoded"]
    wav_path: Path
    prompt_frames: int | None = None
    num_frames: int | None = None

    def __post_init__(self) -> None:
        if self.kind == "decoded" and (self.prompt_frames is None or self.num_frames is None):
            raise ValueError("a decoded reference needs prompt_frames and num_frames")


@dataclass(frozen=True)
class PromptListEntry:
    """One line of a ``--prompt-list`` file with the id ``nida say`` gave the prompt."""

    prompt_id: str
    wav_path: Path
    transcript: str


@dataclass(frozen=True)
class ExternalPromptInfo:
    """What ``settings.json`` of a say run records about one external prompt.

    ``source`` is the path as given to ``--prompt-list``; ``text`` (the normalised transcript)
    and ``sha256`` (of the recording when ``say`` ran) are ``None`` for runs from before they
    were recorded, or if the file was not readable then.
    """

    source: str
    text: str | None = None
    sha256: str | None = None


def chunk_id(prompt_id: str, index: int) -> str:
    """Chunk name used by ``say`` for its files: ``<prompt_id>__<index:02d>``."""
    return f"{prompt_id}__{index:02d}"


def _optional_count(record: Mapping[str, object], key: str, where: str) -> int | None:
    value = record.get(key)
    if value is None:
        return None
    if isinstance(value, bool) or not isinstance(value, int) or value < 0:
        raise SpeakerSimInputError(f"{where}: bad {key} {value!r}")
    return value


def read_say_chunks(say_dir: Path) -> list[ChunkRef]:
    """The chunks listed in ``<say_dir>/say.jsonl`` with their 16 kHz wav paths.

    Raises:
        SpeakerSimInputError: missing / empty / unparseable manifest, a record without a
            valid ``prompt_id`` or ``index`` (or a malformed ``prompt_frames`` /
            ``num_frames``), a duplicate chunk, or missing chunk audio (the first missing
            path is named).
    """
    path = say_dir / SAY_MANIFEST
    if not path.is_file():
        raise SpeakerSimInputError(f"{path}: no say manifest (a `nida say` output directory?)")
    refs: list[ChunkRef] = []
    seen: set[str] = set()
    for line_no, line in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1):
        if not line.strip():
            continue
        where = f"{path} line {line_no}"
        try:
            record = json.loads(line)
        except json.JSONDecodeError as exc:
            raise SpeakerSimInputError(f"{where}: invalid JSON ({exc})") from exc
        if not isinstance(record, dict):
            raise SpeakerSimInputError(f"{where}: expected a JSON object")
        prompt_id, index = record.get("prompt_id"), record.get("index")
        if not isinstance(prompt_id, str) or not prompt_id:
            raise SpeakerSimInputError(f"{where}: bad prompt_id {prompt_id!r}")
        if isinstance(index, bool) or not isinstance(index, int) or index < 0:
            raise SpeakerSimInputError(f"{where}: bad index {index!r}")
        cid = chunk_id(prompt_id, index)
        if cid in seen:
            raise SpeakerSimInputError(f"{where}: duplicate chunk {cid!r}")
        seen.add(cid)
        refs.append(
            ChunkRef(
                cid,
                prompt_id,
                index,
                say_dir / "chunks" / f"{cid}.16k.wav",
                _optional_count(record, "prompt_frames", where),
                _optional_count(record, "num_frames", where),
            )
        )
    if not refs:
        raise SpeakerSimInputError(f"{path}: no chunks")
    missing = [str(r.wav_path) for r in refs if not r.wav_path.is_file()]
    if missing:
        raise SpeakerSimInputError(
            f"{len(missing)} chunk audio file(s) missing, first: {missing[0]} "
            "(`say` writes chunks/<chunk>.16k.wav when it decodes the audio)"
        )
    return refs


def _read_settings(say_dir: Path) -> dict[str, object]:
    path = say_dir / SAY_SETTINGS
    if not path.is_file():
        return {}
    try:
        settings = json.loads(path.read_text(encoding="utf-8"))
    except json.JSONDecodeError as exc:
        raise SpeakerSimInputError(f"{path}: invalid JSON ({exc})") from exc
    if not isinstance(settings, dict):
        raise SpeakerSimInputError(f"{path}: expected a JSON object")
    return settings


def _string_map(settings: Mapping[str, object], key: str, where: Path) -> dict[str, str]:
    value = settings.get(key, {})
    if not isinstance(value, dict) or not all(
        isinstance(k, str) and isinstance(v, str) for k, v in value.items()
    ):
        raise SpeakerSimInputError(f"{where}: {key!r} must map ids to strings")
    return dict(value)


def read_external_prompts(say_dir: Path) -> dict[str, ExternalPromptInfo]:
    """The external prompts recorded in ``settings.json`` (``{}`` if the file or the key is
    absent: a run with dataset prompts only, or an older run).

    Raises:
        SpeakerSimInputError: unparseable settings, a malformed mapping, a transcript or hash
            for an id that is not an external prompt, or a hash that is not 64 hex digits.
    """
    path = say_dir / SAY_SETTINGS
    settings = _read_settings(say_dir)
    sources = _string_map(settings, "external_prompts", path)
    texts = _string_map(settings, "external_prompt_texts", path)
    digests = _string_map(settings, "external_prompt_sha256", path)
    stray = sorted((set(texts) | set(digests)) - set(sources))
    if stray:
        raise SpeakerSimInputError(
            f"{path}: transcript / hash for unknown external prompts {stray}"
        )
    for pid, digest in digests.items():
        if len(digest) != 64 or set(digest) - set("0123456789abcdef"):
            raise SpeakerSimInputError(f"{path}: external_prompt_sha256[{pid!r}] is not a sha256")
    return {
        pid: ExternalPromptInfo(src, texts.get(pid), digests.get(pid))
        for pid, src in sources.items()
    }


def read_say_context(say_dir: Path) -> dict[str, object]:
    """Checkpoint, step, ``w_emo`` and seed of the say run, when ``settings.json`` has them."""
    settings = _read_settings(say_dir)
    return {k: settings[k] for k in ("checkpoint", "step", "w_emo", "seed") if k in settings}


def prompt_id_for(wav: Path, text: str) -> str:
    """The id ``nida say`` gives an external prompt: ``ext-`` + 12 hex of a sha1 over the
    resolved wav path and the (normalised) transcript. A copy of
    ``nida.infer.external_prompt._prompt_id`` (that module needs torch; a test pins equality).

    The id depends on the path: a moved recording gets another id, so prompts are matched to
    recordings by transcript and content hash (:func:`resolve_references`), not by this id alone.
    """
    digest = hashlib.sha1(f"{wav.resolve()}\x1f{text}".encode()).hexdigest()
    return f"ext-{digest[:12]}"


def read_prompt_list(path: Path) -> list[PromptListEntry]:
    """``<wav>\\t<transcript>`` lines (blank lines and ``#`` comments skipped), with prompt ids.

    Same format and normalisation as ``nida.infer.external_prompt.read_prompt_list`` (a test
    pins equality of the ids).

    Raises:
        SpeakerSimInputError: a line without a tab, an empty transcript, a missing wav, an
            empty list, or two lines with the same prompt id.
    """
    entries: list[PromptListEntry] = []
    for number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1):
        if not line.strip() or line.lstrip().startswith("#"):
            continue
        wav, sep, text = line.partition("\t")
        if not sep or not text.strip():
            raise SpeakerSimInputError(f"{path}:{number}: expected '<wav>\\t<transcript>'")
        wav_path = Path(wav.strip())
        if not wav_path.is_file():
            raise SpeakerSimInputError(f"{path}:{number}: {wav_path} does not exist")
        transcript = " ".join(nfc(text).split())
        entries.append(PromptListEntry(prompt_id_for(wav_path, transcript), wav_path, transcript))
    if not entries:
        raise SpeakerSimInputError(f"{path}: no prompt lines")
    ids = [e.prompt_id for e in entries]
    if len(set(ids)) != len(ids):
        raise SpeakerSimInputError(f"{path}: the same (wav, transcript) is listed twice")
    return entries


def _match_prompt_list(
    pid: str, info: ExternalPromptInfo | None, prompt_list: Sequence[PromptListEntry]
) -> int | None:
    """Index of the prompt-list entry for external prompt ``pid``, or ``None``.

    By id first (the recording did not move); else by transcript, narrowed by the recorded
    content hash when several lines share the transcript.

    Raises:
        SpeakerSimInputError: ambiguous transcript match, or a matched file whose content
            differs from the recording ``say`` used.
    """
    found = [i for i, e in enumerate(prompt_list) if e.prompt_id == pid]
    if not found and info is not None and info.text is not None:
        found = [i for i, e in enumerate(prompt_list) if e.transcript == info.text]
        if len(found) > 1 and info.sha256 is not None:
            found = [i for i in found if sha256_file(prompt_list[i].wav_path) == info.sha256]
        if len(found) > 1:
            raise SpeakerSimInputError(
                f"prompt {pid}: {len(found)} prompt-list lines have its transcript and cannot "
                "be told apart (no recorded recording hash)"
            )
    if not found:
        return None
    index = found[0]
    entry = prompt_list[index]
    if info is not None and info.sha256 is not None and sha256_file(entry.wav_path) != info.sha256:
        raise SpeakerSimInputError(
            f"prompt {pid}: {entry.wav_path} differs from the recording `say` used (sha256)"
        )
    return index


def _source_for(
    pid: str,
    info: ExternalPromptInfo | None,
    prompt_list: Sequence[PromptListEntry],
    matched: set[int],
) -> Path:
    index = _match_prompt_list(pid, info, prompt_list)
    if index is not None:
        matched.add(index)
        return prompt_list[index].wav_path
    if info is None:
        raise SpeakerSimInputError(
            f"prompt {pid}: an external prompt whose source recording is not recorded in "
            "settings.json; pass the --prompt-list used for the run, or --reference decoded"
        )
    source = Path(info.source)
    if not source.is_absolute():
        raise SpeakerSimInputError(
            f"prompt {pid}: the recorded source path {info.source!r} is relative, so it cannot "
            "be found reliably; pass the --prompt-list used for the run, or --reference decoded"
        )
    if not source.is_file():
        raise SpeakerSimInputError(
            f"prompt {pid}: source recording {source} does not exist (pass the --prompt-list "
            "used for the run, or --reference decoded)"
        )
    if info.sha256 is not None and sha256_file(source) != info.sha256:
        raise SpeakerSimInputError(
            f"prompt {pid}: {source} differs from the recording `say` used (sha256)"
        )
    return source


def _decoded_reference(say_dir: Path, pid: str, chunks: Sequence[ChunkRef]) -> ReferenceRef:
    first = min(chunks, key=lambda c: c.index)
    if first.prompt_frames is None or first.num_frames is None:
        raise SpeakerSimInputError(
            f"prompt {pid}: say.jsonl has no prompt_frames / num_frames (an older run), so the "
            "decoded prompt cannot be located; use --reference source"
        )
    if not 0 < first.prompt_frames <= first.num_frames:
        raise SpeakerSimInputError(
            f"prompt {pid}: bad frame counts {first.prompt_frames} / {first.num_frames}"
        )
    audio = say_dir / "audio" / f"{first.chunk_id}.16k.wav"
    if not audio.is_file():
        raise SpeakerSimInputError(
            f"prompt {pid}: decoded audio {audio} is missing (`say` writes audio/<chunk>.16k.wav)"
        )
    return ReferenceRef("decoded", audio, first.prompt_frames, first.num_frames)


def resolve_references(
    say_dir: Path,
    chunks: Sequence[ChunkRef],
    *,
    mode: ReferenceMode,
    external: Mapping[str, ExternalPromptInfo],
    prompt_list: Sequence[PromptListEntry] = (),
) -> tuple[dict[str, ReferenceRef], set[int]]:
    """The reference audio of every prompt of a say run, and the prompt-list lines it used.

    An *external* prompt is one recorded in ``settings.json`` or with an ``ext-`` id. In
    ``auto`` / ``source`` mode it is compared with its **source recording**, found in
    ``prompt_list`` (by id, else by transcript, narrowed and verified by the recorded content
    hash) or at the absolute path recorded in ``settings.json`` (verified by the hash when
    recorded). A dataset prompt has no recording in the say directory: ``auto`` compares it with
    the **decoded prompt** (the prompt part of ``audio/<chunk>.16k.wav``, cut like the chunks),
    ``source`` refuses it. ``decoded`` uses the decoded prompt for every prompt. An external
    prompt whose source cannot be found is an error in ``auto`` and ``source`` mode, never a
    silent fall-back to the decoded prompt.

    Returns:
        ``(prompt id -> reference, indices into prompt_list that matched a prompt)``.

    Raises:
        SpeakerSimInputError: a needed file or frame count is missing, a recorded path is
            relative, a hash differs, a transcript match is ambiguous, or ``source`` mode meets a
            dataset prompt.
    """
    by_prompt: dict[str, list[ChunkRef]] = defaultdict(list)
    for c in chunks:
        by_prompt[c.prompt_id].append(c)
    refs: dict[str, ReferenceRef] = {}
    matched: set[int] = set()
    no_source: list[str] = []
    for pid in sorted(by_prompt):
        info = external.get(pid)
        is_external = info is not None or pid.startswith("ext-")
        if mode != "decoded" and is_external:
            refs[pid] = ReferenceRef("source", _source_for(pid, info, prompt_list, matched))
        elif mode == "source":
            no_source.append(pid)
        else:
            refs[pid] = _decoded_reference(say_dir, pid, by_prompt[pid])
    if no_source:
        raise SpeakerSimInputError(
            f"--reference source: no source recording known for prompt(s) {no_source} "
            "(dataset prompts have none; use --reference auto or decoded)"
        )
    return refs, matched


def unmatched_prompt_entries(
    prompt_list: Sequence[PromptListEntry], matched: set[int]
) -> list[PromptListEntry]:
    """The prompt-list lines that no prompt of any say run matched (``matched`` is the union of
    the index sets :func:`resolve_references` returned): a stale or wrong list."""
    return [e for i, e in enumerate(prompt_list) if i not in matched]


# ----------------------------------------------------------------------------- say: scoring


def score_say(
    chunks: Sequence[ChunkRef],
    references: Mapping[str, ReferenceRef],
    reference_outcomes: Mapping[str, EmbedOutcome],
    chunk_outcomes: Mapping[str, EmbedOutcome],
    threshold: float,
) -> list[dict[str, object]]:
    """One record per chunk: similarity to its prompt's reference, or the reason it has none.

    ``reference_outcomes`` is keyed by prompt id, ``chunk_outcomes`` by chunk id. A record
    is ``{chunk, prompt_id, index, reference_kind, similarity, below_threshold, duration_ms}``
    or ``{chunk, prompt_id, index, reference_kind, error}`` (the chunk's own failure, or the
    failure of its reference). ``below_threshold`` is ``similarity < threshold``.

    Raises:
        SpeakerSimInputError: an outcome is missing for a chunk or prompt.
    """
    records: list[dict[str, object]] = []
    for chunk in chunks:
        head: dict[str, object] = {
            "chunk": chunk.chunk_id,
            "prompt_id": chunk.prompt_id,
            "index": chunk.index,
            "reference_kind": references[chunk.prompt_id].kind,
        }
        ref = reference_outcomes.get(chunk.prompt_id)
        out = chunk_outcomes.get(chunk.chunk_id)
        if ref is None or out is None:
            raise SpeakerSimInputError(f"{chunk.chunk_id}: no embedding outcome")
        if isinstance(ref, ClipFailure):
            records.append({**head, "error": f"reference embedding failed: {ref.error}"})
        elif isinstance(out, ClipFailure):
            records.append({**head, "error": out.error})
        else:
            sim = cosine(ref.vector, out.vector)
            records.append(
                {
                    **head,
                    "similarity": sim,
                    "below_threshold": sim < threshold,
                    "duration_ms": out.duration_ms,
                }
            )
    return records


def _similarity_stats(
    records: Sequence[Mapping[str, object]], threshold: float
) -> dict[str, object]:
    sims = [float(str(r["similarity"])) for r in records if "similarity" in r]
    stats = describe(sims)
    below = sum(1 for s in sims if s < threshold)
    return {
        "n_scored": len(sims),
        "mean": stats["mean"],
        "min": stats["min"],
        "max": stats["max"],
        "std": stats["std"],
        "spread": (max(sims) - min(sims)) if sims else None,
        "n_below_threshold": below,
        "drift_fraction": (below / len(sims)) if sims else None,
    }


def summarize_say(
    records: Sequence[Mapping[str, object]],
    references: Mapping[str, ReferenceRef],
    reference_outcomes: Mapping[str, EmbedOutcome],
    threshold: float,
) -> dict[str, object]:
    """Per-reference-kind and per-prompt statistics of :func:`score_say` records.

    ``overall_by_reference_kind`` maps ``"source"`` / ``"decoded"`` (the kinds that occur) to the
    statistics of the chunks compared with that kind of reference. They are never pooled into one
    number: a similarity to the original recording and one to the VAE-decoded prompt are not on the
    same scale. Those entries and every ``per_prompt`` entry hold: ``n_scored``; ``mean`` /
    ``min`` / ``max`` / ``std`` of the similarity; ``spread`` (max - min across the chunks: the
    within-paragraph spread); ``n_below_threshold`` and ``drift_fraction`` (the share of scored
    chunks below ``threshold``). A kind adds ``n_chunks`` and ``n_failed``; a prompt adds
    ``reference_kind``, ``reference_path``, ``n_chunks``, ``n_failed`` and
    ``reference_duration_ms`` (``None`` if its reference failed). Failed chunks are listed with
    their error, never dropped; ``lowest_chunks`` are the weakest few (with their kind).
    """
    per_prompt: dict[str, dict[str, object]] = {}
    by_prompt: dict[str, list[Mapping[str, object]]] = defaultdict(list)
    by_kind: dict[str, list[Mapping[str, object]]] = defaultdict(list)
    for record in records:
        by_prompt[str(record["prompt_id"])].append(record)
        by_kind[str(record["reference_kind"])].append(record)
    for pid in sorted(by_prompt):
        mine = by_prompt[pid]
        ref = reference_outcomes.get(pid)
        per_prompt[pid] = {
            "reference_kind": references[pid].kind,
            "reference_path": str(references[pid].wav_path),
            "reference_duration_ms": ref.duration_ms if isinstance(ref, ClipEmbedding) else None,
            "n_chunks": len(mine),
            "n_failed": sum(1 for r in mine if "error" in r),
            **_similarity_stats(mine, threshold),
        }
    overall = {
        kind: {
            "n_chunks": len(mine),
            "n_failed": sum(1 for r in mine if "error" in r),
            **_similarity_stats(mine, threshold),
        }
        for kind, mine in sorted(by_kind.items())
    }
    failed = [{"chunk": r["chunk"], "error": r["error"]} for r in records if "error" in r]
    scored = sorted(
        (r for r in records if "similarity" in r), key=lambda r: float(str(r["similarity"]))
    )
    return {
        "threshold": threshold,
        "n_chunks": len(records),
        "n_failed": len(failed),
        "failed_chunks": failed,
        "overall_by_reference_kind": overall,
        "per_prompt": per_prompt,
        "lowest_chunks": [
            {
                "chunk": r["chunk"],
                "similarity": r["similarity"],
                "reference_kind": r["reference_kind"],
            }
            for r in scored[:LOWEST_CHUNKS]
        ],
    }


# ---------------------------------------------------------------------------- groups: input


@dataclass(frozen=True)
class IndexRow:
    """The fields of a latent dataset ``index.jsonl`` row that the group check needs."""

    segment_id: str
    group: str
    split: str
    num_frames: int


def read_index(dataset_dir: Path) -> tuple[list[IndexRow], float]:
    """Rows of ``<dataset_dir>/index.jsonl`` and the latent frame hop (ms) from ``meta.json``.

    Raises:
        SpeakerSimInputError: a missing / empty / malformed index, a duplicate segment id, or
            a ``meta.json`` without a positive ``frame_hop_ms``.
    """
    index_path, meta_path = dataset_dir / "index.jsonl", dataset_dir / "meta.json"
    if not index_path.is_file() or not meta_path.is_file():
        raise SpeakerSimInputError(
            f"{dataset_dir}: needs index.jsonl and meta.json (a latent dataset directory)"
        )
    try:
        meta = json.loads(meta_path.read_text(encoding="utf-8"))
    except json.JSONDecodeError as exc:
        raise SpeakerSimInputError(f"{meta_path}: invalid JSON ({exc})") from exc
    hop = meta.get("frame_hop_ms") if isinstance(meta, dict) else None
    if isinstance(hop, bool) or not isinstance(hop, int | float) or not hop > 0:
        raise SpeakerSimInputError(f"{meta_path}: frame_hop_ms must be a positive number")
    rows: list[IndexRow] = []
    seen: set[str] = set()
    for line_no, line in enumerate(index_path.read_text(encoding="utf-8").splitlines(), start=1):
        if not line.strip():
            continue
        where = f"{index_path} line {line_no}"
        try:
            record = json.loads(line)
        except json.JSONDecodeError as exc:
            raise SpeakerSimInputError(f"{where}: invalid JSON ({exc})") from exc
        if not isinstance(record, dict):
            raise SpeakerSimInputError(f"{where}: expected a JSON object")
        sid, group, split = record.get("segment_id"), record.get("group"), record.get("split")
        frames = record.get("num_frames")
        if not isinstance(sid, str) or not sid or not isinstance(group, str) or not group:
            raise SpeakerSimInputError(f"{where}: bad segment_id / group")
        if not isinstance(split, str) or isinstance(frames, bool) or not isinstance(frames, int):
            raise SpeakerSimInputError(f"{where}: bad split / num_frames")
        if sid in seen:
            raise SpeakerSimInputError(f"{where}: duplicate segment {sid!r}")
        seen.add(sid)
        rows.append(IndexRow(sid, group, split, frames))
    if not rows:
        raise SpeakerSimInputError(f"{index_path}: no rows")
    return rows, float(hop)


def find_dataset_dir(root: Path) -> Path:
    """The one finished ``dataset-*`` directory (with ``meta.json``, written last) of an importer
    root.

    Raises:
        SpeakerSimInputError: none, or more than one (pass the directory explicitly).
    """
    found = sorted(p for p in root.glob("dataset-*") if (p / "meta.json").is_file())
    if len(found) != 1:
        raise SpeakerSimInputError(
            f"{root}: expected exactly one finished dataset-*/ directory, found "
            f"{[p.name for p in found]} (use --dataset-dir)"
        )
    return found[0]


def list_audio_stems(audio_root: Path, suffix: str = ".wav") -> set[str]:
    """File stems (segment ids) of ``<audio_root>/*<suffix>``, from one directory listing.

    Raises:
        SpeakerSimInputError: if ``audio_root`` is not a directory.
    """
    if not audio_root.is_dir():
        raise SpeakerSimInputError(
            f"{audio_root}: not a directory (the importer's 16 kHz wavs live in its local "
            "cache, <cache>/audio16k; rerun `nida data-import --stage audio` to rebuild them)"
        )
    with os.scandir(audio_root) as entries:
        return {e.name[: -len(suffix)] for e in entries if e.name.endswith(suffix) and e.is_file()}


def select_candidates(
    rows: Sequence[IndexRow],
    available: set[str],
    *,
    hop_ms: float,
    min_duration_ms: float,
    split: str | None = None,
) -> tuple[dict[str, list[str]], dict[str, int]]:
    """Segment ids per group that can be embedded, and why the others cannot.

    A row is dropped (and counted, never silently) for: another ``split`` than requested,
    ``num_frames * hop_ms < min_duration_ms`` (``too_short``: an embedding of a very short
    clip is unreliable), or no wav in ``available`` (``missing_audio``).

    Returns:
        ``(group -> sorted segment ids, counts)`` with the counts ``n_rows``, ``n_other_split``,
        ``n_too_short``, ``n_missing_audio`` and ``n_candidates``.
    """
    counts = {
        "n_rows": len(rows),
        "n_other_split": 0,
        "n_too_short": 0,
        "n_missing_audio": 0,
        "n_candidates": 0,
    }
    groups: dict[str, list[str]] = defaultdict(list)
    for row in rows:
        if split is not None and row.split != split:
            counts["n_other_split"] += 1
        elif row.num_frames * hop_ms < min_duration_ms:
            counts["n_too_short"] += 1
        elif row.segment_id not in available:
            counts["n_missing_audio"] += 1
        else:
            groups[row.group].append(row.segment_id)
            counts["n_candidates"] += 1
    return {g: sorted(ids) for g, ids in sorted(groups.items())}, counts


def sample_groups(
    candidates: Mapping[str, Sequence[str]], *, n_groups: int, per_group: int, seed: int
) -> list[tuple[str, tuple[str, ...]]]:
    """``n_groups`` groups with at least ``per_group`` candidates, ``per_group`` segments each.

    Deterministic: a pure function of the candidate sets, ``seed`` and the arguments (groups and
    segments are sorted before sampling). Fewer eligible groups than ``n_groups`` yields all of
    them. The result is ordered by group id, segments by id.

    Raises:
        ValueError: ``n_groups < 1`` or ``per_group < 2`` (one segment has no intra pair).
    """
    if n_groups < 1:
        raise ValueError(f"n_groups must be >= 1, got {n_groups}")
    if per_group < 2:
        raise ValueError(f"per_group must be >= 2, got {per_group}")
    rng = random.Random(seed)
    eligible = sorted(g for g, ids in candidates.items() if len(set(ids)) >= per_group)
    chosen = sorted(rng.sample(eligible, min(n_groups, len(eligible))))
    return [(g, tuple(sorted(rng.sample(sorted(set(candidates[g])), per_group)))) for g in chosen]


# --------------------------------------------------------------------------- groups: scoring


def _pair_similarities(vectors: Sequence[FloatVec]) -> list[float]:
    matrix = np.stack(vectors).astype(np.float64)
    gram = matrix @ matrix.T
    upper = np.triu_indices(len(vectors), k=1)
    return [float(x) for x in gram[upper]]


def _auc(positives: Sequence[float], negatives: Sequence[float]) -> float:
    """P(positive > negative) with ties counting half (the area under the ROC curve)."""
    neg = np.sort(np.asarray(negatives, dtype=np.float64))
    pos = np.asarray(positives, dtype=np.float64)
    below = np.searchsorted(neg, pos, side="left")
    at_most = np.searchsorted(neg, pos, side="right")
    return float(np.mean((below + 0.5 * (at_most - below)) / neg.size))


def _inter_similarities(
    embedded: Mapping[str, Sequence[tuple[str, FloatVec]]], pairs: int, seed: int
) -> list[float]:
    """Cosine of ``pairs`` random segment pairs drawn from two different groups."""
    groups = sorted(embedded)
    if len(groups) < 2:
        return []
    rng = random.Random(seed)
    sims: list[float] = []
    for _ in range(pairs):
        first, second = rng.sample(groups, 2)
        a = rng.choice(embedded[first])[1].astype(np.float64)
        b = rng.choice(embedded[second])[1].astype(np.float64)
        sims.append(float(np.dot(a, b)))
    return sims


def summarize_groups(
    samples: Sequence[tuple[str, tuple[str, ...]]],
    outcomes: Mapping[str, EmbedOutcome],
    *,
    intra_threshold: float,
    inter_pairs: int,
    seed: int,
) -> tuple[list[dict[str, object]], dict[str, object]]:
    """Per-group records and the intra- vs inter-group summary of the sampled segments.

    Per group (``outcomes`` is keyed by segment id): ``group``, ``segments``, ``n_embedded``,
    ``errors`` (segment -> error, never dropped) and, when at least two segments were embedded,
    ``intra_mean`` / ``intra_min`` (cosine over all pairs) and ``below_threshold``
    (``intra_mean < intra_threshold``).

    Summary: ``intra.pairs`` / ``intra.group_means`` / ``inter`` (:func:`describe`), the counts,
    ``n_groups_below`` / ``fraction_groups_below`` (over groups with an intra mean),
    ``separation`` (mean intra pair minus mean inter pair) and ``auc`` (probability that a random
    intra pair is more similar than a random inter pair). ``inter`` uses ``inter_pairs``
    random pairs of embedded segments from two different groups (seeded; empty stats if fewer
    than two groups have an embedded segment).

    Raises:
        ValueError: ``inter_pairs < 1``.
        SpeakerSimInputError: a sampled segment has no outcome.
    """
    if inter_pairs < 1:
        raise ValueError(f"inter_pairs must be >= 1, got {inter_pairs}")
    records: list[dict[str, object]] = []
    embedded: dict[str, list[tuple[str, FloatVec]]] = {}
    intra_pair_sims: list[float] = []
    group_means: list[float] = []
    n_failed = 0
    for group, segments in samples:
        good: list[tuple[str, FloatVec]] = []
        errors: dict[str, str] = {}
        for sid in segments:
            outcome = outcomes.get(sid)
            if outcome is None:
                raise SpeakerSimInputError(f"{sid}: no embedding outcome")
            if isinstance(outcome, ClipFailure):
                errors[sid] = outcome.error
            else:
                good.append((sid, outcome.vector))
        n_failed += len(errors)
        record: dict[str, object] = {
            "group": group,
            "segments": list(segments),
            "n_embedded": len(good),
            "errors": errors,
        }
        if good:
            embedded[group] = good
        if len(good) >= 2:
            sims = _pair_similarities([v for _, v in good])
            mean = float(np.mean(sims))
            record["intra_mean"] = mean
            record["intra_min"] = min(sims)
            record["below_threshold"] = mean < intra_threshold
            intra_pair_sims.extend(sims)
            group_means.append(mean)
        records.append(record)
    inter_sims = _inter_similarities(embedded, inter_pairs, seed)
    below = sum(1 for m in group_means if m < intra_threshold)
    intra_mean_all = float(np.mean(intra_pair_sims)) if intra_pair_sims else None
    inter_mean_all = float(np.mean(inter_sims)) if inter_sims else None
    summary: dict[str, object] = {
        "n_groups_sampled": len(samples),
        "n_groups_with_intra": len(group_means),
        "n_segments_failed": n_failed,
        "intra_threshold": intra_threshold,
        "n_groups_below": below,
        "fraction_groups_below": (below / len(group_means)) if group_means else None,
        "intra": {"pairs": describe(intra_pair_sims), "group_means": describe(group_means)},
        "inter": describe(inter_sims),
        "separation": (
            intra_mean_all - inter_mean_all
            if intra_mean_all is not None and inter_mean_all is not None
            else None
        ),
        "auc": _auc(intra_pair_sims, inter_sims) if intra_pair_sims and inter_sims else None,
    }
    return records, summary