"""S10: group-disjoint train/val/test splits (docs/03-data-pipeline.md).

Whole *groups* (speakers) are assigned to a split, never single items, so a speaker heard in
training is never in validation or test. The assignment is a pure function of
``(seed, group_id)``: it does not depend on which other groups exist, on item order, or on
how many items a group has. Adding data later therefore never moves an already assigned
group to another split.

What a "group" is must be chosen so that it really identifies a speaker. For YODAS the
diarization labels (``SPEAKER_00`` ...) are only unique *within one file*, so the group is
the source video/channel id, not the diarization label: two videos of the same channel may
share a speaker, and a channel is the finest unit that cannot leak one speaker across
splits. For corpora with global speaker ids, use the speaker id.

Group ids are normalised to NFC before hashing and comparison (two visually identical
"ş" must land in the same split). Item ids are used as given.

Because whole groups move together, realised split sizes deviate from the requested
fractions when there are few groups or very uneven group sizes; the caller should report
the realised sizes.
"""

from __future__ import annotations

import hashlib
import math
from collections.abc import Mapping

from nida.text.turkish import nfc

__all__ = ["SplitError", "assign_splits", "check_disjoint", "group_uniform"]

_FRACTION_SUM_TOLERANCE = 1e-9
_HASH_BITS = 53  # a double has a 53-bit mantissa: top-53-bits / 2**53 is exact and < 1
_DIGEST_BYTES = 8


class SplitError(ValueError):
    """Invalid split request or a split assignment that violates group disjointness."""


def group_uniform(group_id: str, seed: str) -> float:
    """Deterministic pseudo-uniform value in ``[0, 1)`` for a group under ``seed``.

    ``sha1(seed U+001F nfc(group_id))``; the separator makes ``("ab", "c")`` and
    ``("a", "bc")`` hash differently.
    """
    key = "\x1f".join((seed, nfc(group_id))).encode("utf-8")
    top = int.from_bytes(hashlib.sha1(key).digest()[:_DIGEST_BYTES], "big")
    return (top >> (_DIGEST_BYTES * 8 - _HASH_BITS)) / float(1 << _HASH_BITS)


def _validate_fractions(fractions: Mapping[str, float]) -> list[tuple[str, float]]:
    if not fractions:
        raise SplitError("fractions must not be empty")
    for name, fraction in fractions.items():
        if not name:
            raise SplitError("split names must be non-empty")
        if not (math.isfinite(fraction) and 0.0 <= fraction <= 1.0):
            raise SplitError(f"fraction for split {name!r} must be in [0, 1], got {fraction!r}")
    total = math.fsum(fractions.values())
    if abs(total - 1.0) > _FRACTION_SUM_TOLERANCE:
        raise SplitError(f"fractions must sum to 1 (tolerance 1e-9), got {total!r}")
    return sorted(fractions.items())  # sorted split names: stable regardless of dict order


def assign_splits(
    groups: Mapping[str, str], fractions: Mapping[str, float], seed: str
) -> dict[str, str]:
    """Assign every item to a split such that each group lands in exactly one split.

    Args:
        groups: ``{item_id: group_id}``.
        fractions: ``{split_name: fraction}``, each in [0, 1], summing to 1. Split names are
            ordered alphabetically for the cumulative boundaries, so the result does not
            depend on mapping order. A split with fraction 0 receives nothing.
        seed: Non-empty reproducibility string; part of the run's recorded config.

    Returns:
        ``{item_id: split_name}`` for every key of ``groups``.

    Raises:
        SplitError: on invalid fractions, an empty seed, or an empty item or group id.
    """
    if not seed:
        raise SplitError("seed must be non-empty")
    ordered = _validate_fractions(fractions)
    boundaries: list[tuple[str, float]] = []
    cumulative = 0.0
    for name, fraction in ordered:
        cumulative += fraction
        boundaries.append((name, cumulative))

    fallback = next(name for name, fraction in reversed(ordered) if fraction > 0.0)
    split_of_group: dict[str, str] = {}
    assignment: dict[str, str] = {}
    for item_id, group_id in groups.items():
        if not item_id:
            raise SplitError("empty item id")
        if not group_id:
            raise SplitError(f"item {item_id!r} has an empty group id")
        key = nfc(group_id)
        split = split_of_group.get(key)
        if split is None:
            value = group_uniform(key, seed)
            # Values >= the last boundary (float rounding of the cumulative sum) fall into
            # the last split with a non-zero share, never into a zero-fraction one.
            split = next((name for name, upper in boundaries if value < upper), fallback)
            split_of_group[key] = split
        assignment[item_id] = split
    return assignment


def check_disjoint(assignment: Mapping[str, str], groups: Mapping[str, str]) -> None:
    """Verify that no group spans more than one split.

    Args:
        assignment: ``{item_id: split_name}``.
        groups: ``{item_id: group_id}``.

    Raises:
        SplitError: if the two mappings do not cover the same items, or a group (compared
            after NFC) has items in more than one split. The message names offending groups.
    """
    only_assigned = sorted(set(assignment) - set(groups))
    only_grouped = sorted(set(groups) - set(assignment))
    if only_assigned or only_grouped:
        raise SplitError(
            f"assignment and groups cover different items "
            f"(no group: {only_assigned[:3]}, no split: {only_grouped[:3]})"
        )
    splits_of_group: dict[str, set[str]] = {}
    for item_id, group_id in groups.items():
        splits_of_group.setdefault(nfc(group_id), set()).add(assignment[item_id])
    leaking = {g: sorted(s) for g, s in sorted(splits_of_group.items()) if len(s) > 1}
    if leaking:
        shown = ", ".join(f"{g!r}: {s}" for g, s in list(leaking.items())[:5])
        raise SplitError(f"{len(leaking)} group(s) span several splits: {shown}")