src/turn_detection/data/splits.py

"""Deterministic, group-safe and approximately stratified split assignment."""

from __future__ import annotations

from collections import Counter, defaultdict
import hashlib
import heapq
import json
import math
from typing import Any, Iterable, Iterator, Mapping, MutableMapping, Sequence

from .grouping import attach_group_ids


DEFAULT_SPLIT_RATIOS: dict[str, float] = {"train": 0.9, "validation": 0.1}


class SplitValidationError(ValueError):
    """Raised when a split manifest contains leakage or invalid assignments."""


def normalize_split_ratios(ratios: Mapping[str, float]) -> dict[str, float]:
    """Validate and normalize positive split weights to sum to one."""

    if not ratios:
        raise ValueError("at least one split ratio is required")
    normalized: dict[str, float] = {}
    for name, raw_value in ratios.items():
        split_name = str(name).strip()
        if not split_name:
            raise ValueError("split names cannot be empty")
        value = float(raw_value)
        if not math.isfinite(value) or value <= 0:
            raise ValueError(f"ratio for {split_name!r} must be finite and positive")
        if split_name in normalized:
            raise ValueError(f"duplicate split name: {split_name}")
        normalized[split_name] = value
    total = sum(normalized.values())
    return {name: value / total for name, value in normalized.items()}


def parse_split_ratios(values: Sequence[str]) -> dict[str, float]:
    """Parse CLI values such as ``train=0.9`` and ``validation=0.1``."""

    parsed: dict[str, float] = {}
    for value in values:
        if "=" not in value:
            raise ValueError(f"split ratio must use NAME=WEIGHT syntax: {value!r}")
        name, raw_ratio = value.split("=", 1)
        name = name.strip()
        if name in parsed:
            raise ValueError(f"duplicate split name: {name}")
        try:
            parsed[name] = float(raw_ratio)
        except ValueError as exc:
            raise ValueError(f"invalid ratio in {value!r}") from exc
    return normalize_split_ratios(parsed)


def _stable_hash(*values: Any) -> int:
    payload = "\0".join(str(value) for value in values).encode("utf-8")
    return int.from_bytes(hashlib.sha256(payload).digest()[:8], "big")


def _category(value: Any) -> str:
    if value is None:
        return "<null>"
    if isinstance(value, bool):
        return "true" if value else "false"
    return str(value).strip() or "<empty>"


def _field_value(row: Mapping[str, Any], field: str) -> Any:
    current: Any = row
    for part in field.split("."):
        if not isinstance(current, Mapping) or part not in current:
            return None
        current = current[part]
    return current


def _add_holdout_links(rows: Sequence[MutableMapping[str, Any]], fields: Sequence[str]) -> None:
    if not fields:
        return
    for row in rows:
        raw_keys = row.get("group_keys") or []
        keys = [str(raw_keys)] if isinstance(raw_keys, str) else [str(key) for key in raw_keys]
        for field in fields:
            value = _field_value(row, field)
            if value is None or value == "":
                continue
            canonical = json.dumps(value, sort_keys=True, ensure_ascii=False, default=str)
            digest = hashlib.sha256(f"{field}\0{canonical}".encode("utf-8")).hexdigest()[:32]
            keys.append(f"holdout:{field}:{digest}")
        row["group_keys"] = sorted(set(keys))
    attach_group_ids(rows)


def assign_splits(
    rows: Sequence[Mapping[str, Any]],
    *,
    ratios: Mapping[str, float] | None = None,
    seed: int = 42,
    stratify_fields: Sequence[str] = ("endpoint", "language", "dataset"),
    holdout_fields: Sequence[str] = (),
) -> list[dict[str, Any]]:
    """Assign whole leakage components using deterministic greedy balancing.

    The objective balances total rows and each requested marginal stratum.  It
    is deterministic for a fixed set of rows regardless of input row order.
    ``holdout_fields`` can enforce source-, speaker-, or conversation-held-out
    evaluation by linking all equal values before assignment.
    """

    split_ratios = normalize_split_ratios(ratios or DEFAULT_SPLIT_RATIOS)
    output = [dict(row) for row in rows]
    if not output:
        return output
    if any(not row.get("group_id") for row in output):
        attach_group_ids(output)
    _add_holdout_links(output, holdout_fields)

    group_indices: dict[str, list[int]] = defaultdict(list)
    for index, row in enumerate(output):
        group_indices[str(row["group_id"])].append(index)

    total_rows = len(output)
    all_strata: Counter[str] = Counter()
    group_strata: dict[str, Counter[str]] = {}
    for group_id, indices in group_indices.items():
        counts: Counter[str] = Counter()
        for index in indices:
            for field in stratify_fields:
                key = f"{field}={_category(_field_value(output[index], field))}"
                counts[key] += 1
                all_strata[key] += 1
        group_strata[group_id] = counts

    # Iterative multilabel stratification: repeatedly satisfy the rarest
    # remaining marginal, assigning one whole group to the split with the
    # greatest remaining quota for that marginal.  Static heaps and lazy
    # deletion keep this O(groups * fields * log(groups)) at full scale.
    remaining_groups = set(group_indices)
    remaining_strata = Counter(all_strata)
    desired_total = {name: total_rows * ratio for name, ratio in split_ratios.items()}
    desired_strata = {
        name: {key: count * ratio for key, count in all_strata.items()}
        for name, ratio in split_ratios.items()
    }
    candidate_heaps: dict[str, list[tuple[int, int, int, str]]] = defaultdict(list)
    fallback_heap: list[tuple[int, int, str]] = []
    for group_id, indices in group_indices.items():
        group_size = len(indices)
        heapq.heappush(
            fallback_heap,
            (-group_size, _stable_hash(seed, "fallback-group", group_id), group_id),
        )
        for stratum, contribution in group_strata[group_id].items():
            heapq.heappush(
                candidate_heaps[stratum],
                (
                    -contribution,
                    -group_size,
                    _stable_hash(seed, "stratum-group", stratum, group_id),
                    group_id,
                ),
            )
    group_assignment: dict[str, str] = {}

    while remaining_groups:
        active_strata = [
            (count, _stable_hash(seed, "stratum-order", stratum), stratum)
            for stratum, count in remaining_strata.items()
            if count > 0
        ]
        focus_stratum = min(active_strata)[2] if active_strata else None
        if focus_stratum is not None:
            candidates = candidate_heaps[focus_stratum]
            while candidates and candidates[0][3] not in remaining_groups:
                heapq.heappop(candidates)
            if not candidates:
                # Defensive fallback for malformed/non-integer stratum counts.
                focus_stratum = None
        if focus_stratum is None:
            while fallback_heap and fallback_heap[0][2] not in remaining_groups:
                heapq.heappop(fallback_heap)
            if not fallback_heap:
                raise RuntimeError("internal split assignment error: no remaining group candidate")
            group_id = fallback_heap[0][2]
        else:
            group_id = candidate_heaps[focus_stratum][0][3]

        group_size = len(group_indices[group_id])
        strata = group_strata[group_id]
        candidate_scores: list[tuple[float, float, float, int, str]] = []
        for split_name in split_ratios:
            primary_need = (
                desired_strata[split_name].get(focus_stratum, 0.0)
                if focus_stratum is not None
                else desired_total[split_name]
            )
            # Secondary need uses all marginals carried by this group.  Dividing
            # by global stratum size prevents rare categories from dominating.
            secondary_need = sum(
                max(desired_strata[split_name].get(stratum, 0.0), 0.0)
                * contribution
                / max(all_strata[stratum], 1)
                for stratum, contribution in strata.items()
            )
            candidate_scores.append(
                (
                    -primary_need,
                    -secondary_need,
                    -desired_total[split_name],
                    _stable_hash(seed, "split-choice", group_id, split_name),
                    split_name,
                )
            )

        selected = min(candidate_scores)[4]
        group_assignment[group_id] = selected
        remaining_groups.remove(group_id)
        desired_total[selected] -= group_size
        for stratum, contribution in strata.items():
            desired_strata[selected][stratum] -= contribution
            remaining_strata[stratum] -= contribution

    for row in output:
        row["split"] = group_assignment[str(row["group_id"])]
    return output


def assign_leave_one_out(
    rows: Sequence[Mapping[str, Any]],
    *,
    field: str,
    held_out_value: Any,
    train_split: str = "train",
    held_out_split: str = "validation",
) -> list[dict[str, Any]]:
    """Assign a named source/domain value to one held-out stress-test split.

    If a pre-existing leakage component contains both held-out and non-held-out
    values, the entire component is held out.  This preserves leakage safety at
    the cost of a small amount of train-domain spillover, which is surfaced by
    :func:`build_split_report`.
    """

    if not field.strip():
        raise ValueError("leave-one-out field cannot be empty")
    if not train_split or not held_out_split or train_split == held_out_split:
        raise ValueError("train and held-out split names must be distinct and non-empty")
    output = [dict(row) for row in rows]
    if any(not row.get("group_id") for row in output):
        attach_group_ids(output)
    expected = _category(held_out_value)
    held_out_groups = {
        str(row["group_id"])
        for row in output
        if _category(_field_value(row, field)) == expected
    }
    if not held_out_groups:
        raise ValueError(f"held-out value {held_out_value!r} was not found in field {field!r}")
    for row in output:
        row["split"] = held_out_split if str(row["group_id"]) in held_out_groups else train_split
    assert_no_split_leakage(output)
    return output


def iter_leave_one_out_folds(
    rows: Sequence[Mapping[str, Any]],
    *,
    field: str = "dataset",
    values: Sequence[Any] | None = None,
    train_split: str = "train",
    held_out_split: str = "validation",
) -> Iterator[tuple[str, list[dict[str, Any]]]]:
    """Yield deterministic leave-one-source/domain-out manifests one at a time."""

    if values is None:
        observed = {
            _category(_field_value(row, field))
            for row in rows
            if _field_value(row, field) is not None and _field_value(row, field) != ""
        }
        selected_values: Sequence[Any] = sorted(observed)
    else:
        selected_values = values
    for value in selected_values:
        canonical = _category(value)
        yield canonical, assign_leave_one_out(
            rows,
            field=field,
            held_out_value=value,
            train_split=train_split,
            held_out_split=held_out_split,
        )


def find_split_leakage(rows: Iterable[Mapping[str, Any]]) -> dict[str, Any]:
    """Detect group, exact-audio and metadata-key crossings between splits."""

    split_by_group: dict[str, set[str]] = defaultdict(set)
    split_by_audio: dict[str, set[str]] = defaultdict(set)
    split_by_key: dict[str, set[str]] = defaultdict(set)
    missing_split = 0
    for row in rows:
        split = row.get("split")
        if not split:
            missing_split += 1
            continue
        split_name = str(split)
        if row.get("group_id"):
            split_by_group[str(row["group_id"])].add(split_name)
        if row.get("audio_sha256"):
            split_by_audio[str(row["audio_sha256"])].add(split_name)
        raw_keys = row.get("group_keys") or []
        if isinstance(raw_keys, str):
            raw_keys = [raw_keys]
        for key in raw_keys:
            split_by_key[str(key)].add(split_name)

    group_crossings = {key: sorted(value) for key, value in split_by_group.items() if len(value) > 1}
    audio_crossings = {key: sorted(value) for key, value in split_by_audio.items() if len(value) > 1}
    key_crossings = {key: sorted(value) for key, value in split_by_key.items() if len(value) > 1}
    return {
        "missing_split_rows": missing_split,
        "group_crossings": group_crossings,
        "audio_hash_crossings": audio_crossings,
        "metadata_key_crossings": key_crossings,
        "is_valid": not (missing_split or group_crossings or audio_crossings or key_crossings),
    }


def assert_no_split_leakage(rows: Iterable[Mapping[str, Any]]) -> None:
    """Raise :class:`SplitValidationError` if any linkage crosses splits."""

    report = find_split_leakage(rows)
    if not report["is_valid"]:
        raise SplitValidationError(
            "split leakage detected: "
            f"missing={report['missing_split_rows']}, groups={len(report['group_crossings'])}, "
            f"audio={len(report['audio_hash_crossings'])}, keys={len(report['metadata_key_crossings'])}"
        )


def build_split_report(
    rows: Sequence[Mapping[str, Any]],
    *,
    stratify_fields: Sequence[str] = ("endpoint", "language", "dataset"),
    holdout_fields: Sequence[str] = (),
) -> dict[str, Any]:
    """Summarize split sizes, strata and leakage validation."""

    counts = Counter(str(row.get("split") or "<missing>") for row in rows)
    by_field: dict[str, dict[str, dict[str, int]]] = {}
    for field in stratify_fields:
        split_values: dict[str, Counter[str]] = defaultdict(Counter)
        for row in rows:
            split_values[str(row.get("split") or "<missing>")][_category(_field_value(row, field))] += 1
        by_field[field] = {
            split: dict(sorted(values.items())) for split, values in sorted(split_values.items())
        }
    holdout_report: dict[str, Any] = {}
    for field in holdout_fields:
        splits_by_value: dict[str, set[str]] = defaultdict(set)
        for row in rows:
            value = _field_value(row, field)
            if value is None or value == "":
                continue
            splits_by_value[_category(value)].add(str(row.get("split") or "<missing>"))
        values_by_split: dict[str, list[str]] = defaultdict(list)
        for value, splits in splits_by_value.items():
            for split in splits:
                values_by_split[split].append(value)
        holdout_report[field] = {
            "by_split": {
                split: sorted(values) for split, values in sorted(values_by_split.items())
            },
            "crossing_values": {
                value: sorted(splits) for value, splits in splits_by_value.items() if len(splits) > 1
            },
        }
    domain_shift: dict[str, Any] | None = None
    if holdout_fields:
        confounder_fields = ("endpoint", "language", "dataset", "synthetic", "midfiller", "endfiller")
        confounder_counts: dict[str, dict[str, dict[str, int]]] = {}
        for field in confounder_fields:
            split_values: dict[str, Counter[str]] = defaultdict(Counter)
            for row in rows:
                split_values[str(row.get("split") or "<missing>")][
                    _category(_field_value(row, field))
                ] += 1
            confounder_counts[field] = {
                split: dict(sorted(values.items()))
                for split, values in sorted(split_values.items())
            }
        domain_shift = {
            "kind": "domain_shift_stress_test",
            "caution": (
                "Held-out source/domain values are confounded with language, synthetic status, "
                "label and collection process; results measure joint domain shift and must not "
                "be interpreted as a clean causal source effect."
            ),
            "slice_counts": confounder_counts,
        }
    return {
        "records": len(rows),
        "split_counts": dict(sorted(counts.items())),
        "strata": by_field,
        "holdouts": holdout_report,
        "domain_shift": domain_shift,
        "leakage": find_split_leakage(rows),
    }