#!/usr/bin/env python3
"""split_rule.py — THE held-out split rule for the house adapter programme.

PREREG-HOUSE-ARMS.md §3.1(4): "The split script is named and lives in one place ...
The freeze script and any dry run import the same module; its sha256 enters the
freeze."  A split rule that exists twice is a split rule that will disagree with
itself, so this file is the only place either rule is written down.

Two rules live here, and they are for two different shapes of corpus:

  * ``s3_artist_split``  — ARM 1 / ARM 2 (`house-wide`, `house-core`).  PREREG
    §3.1(3), candidate **S3**: whole artists, smallest catalogue first,
    DETERMINISTIC, NO RNG.  Registered verbatim as::

        size[a] = |tracks of artist a|                        # over the FROZEN manifest
        order   = sorted(artists, key=lambda a: (size[a], a)) # smallest catalogue first, tie by artist_id
        hold, cnt = [], 0
        for a in order:
            if cnt >= 0.20 * N: break                         # STOP AT THE TARGET
            hold.append(a); cnt += size[a]                    # whole artists only
        # tracks within an artist ordered by track_id; no RNG, no seed, no second pass

  * ``by_track_split``   — ARM 3 (`house-artist`), where the corpus is a single
    artist and "by album" is infeasible (PREREG §3.3).  ``numpy.random
    .default_rng(42).permutation`` over track ids sorted ascending, taking tracks
    until the count first reaches ``ceil(0.20 * N)``.  Seed 42 is fixed by the
    prereg and is the ONLY use of an RNG anywhere in the programme's splits.

numpy is imported **inside** ``by_track_split``, not at module scope, and its
version is reported by ``pinned_stack()`` so that arm 3's permutation carries the
instrument that produced it.  The import site is the proof of the registration:
arm 1's rule is RNG-free, so arm 1's freeze runs on a box with no numpy at all
and cannot silently acquire a random number.  Arm 3's freeze must run where numpy
is installed, and ``pinned_stack()`` is what pins it there.

Nothing in this module reads a file, a clock or an environment variable: it is a
pure function of the manifest rows it is handed.
"""

from __future__ import annotations

import math
from dataclasses import dataclass, field
from typing import Iterable, Mapping, Sequence

# The registered target and the accepted band (PREREG §3.1(2)).
HOLDOUT_TARGET = 0.20
BAND_LOW = 0.18
BAND_HIGH = 0.23

# The arm-3 intra-artist seed (PREREG §3.3).  Arm 1 never consumes it.
BY_TRACK_SEED = 42

RULE_VERBATIM_S3 = (
    "S3 — whole artists, smallest catalogue first, deterministic, no RNG. "
    "size[a] = |tracks of artist a| over the FROZEN manifest; "
    "order = sorted(artists, key=lambda a: (size[a], a)) (smallest catalogue "
    "first, tie broken by artist_id ascending); walk the order appending whole "
    "artists to the holdout while cnt < 0.20 * N, stopping at the first artist "
    "that would be reached with cnt already at or above the target; tracks "
    "within an artist ordered by track_id ascending; no RNG, no seed, no second "
    "pass. Held-out artists are all-in or all-out — no artist straddles the "
    "boundary."
)

RULE_VERBATIM_BY_TRACK = (
    "BY TRACK — for an arm whose corpus is a single artist, where the only "
    "by-album partitions miss the [18 %, 23 %] band. "
    "numpy.random.default_rng(42).permutation over track ids sorted ascending, "
    "taking tracks until the count first reaches ceil(0.20 * N)."
)


@dataclass(frozen=True)
class SplitResult:
    """The outcome of a split, with everything the freeze has to record."""

    rule: str
    rule_verbatim: str
    unit_kind: str
    target_fraction: float
    n_total: int
    holdout_track_ids: tuple[str, ...]
    training_track_ids: tuple[str, ...]
    holdout_units: tuple[str, ...]
    training_units: tuple[str, ...]
    walk: tuple[dict, ...] = field(default=())

    @property
    def n_holdout(self) -> int:
        return len(self.holdout_track_ids)

    @property
    def n_train(self) -> int:
        return len(self.training_track_ids)

    @property
    def actual_fraction(self) -> float:
        return self.n_holdout / self.n_total if self.n_total else 0.0

    @property
    def in_band(self) -> bool:
        return BAND_LOW <= self.actual_fraction <= BAND_HIGH

    def straddlers(self, unit_of: Mapping[str, str]) -> tuple[str, ...]:
        """Units appearing on BOTH sides — the leak D-20260831-03 exists to stop.

        For an artist split this must be empty by construction; it is computed
        rather than asserted so the freeze records a measurement, not a belief.
        """
        held = {unit_of[t] for t in self.holdout_track_ids}
        trained = {unit_of[t] for t in self.training_track_ids}
        return tuple(sorted(held & trained))


def _index_by_unit(
    rows: Iterable[Mapping], unit_key: str, track_key: str
) -> dict[str, list[str]]:
    """Group track ids by their split unit, tracks sorted ascending within a unit."""
    groups: dict[str, list[str]] = {}
    for row in rows:
        groups.setdefault(row[unit_key], []).append(row[track_key])
    for tracks in groups.values():
        tracks.sort()
    return groups


def s3_artist_split(
    rows: Sequence[Mapping],
    *,
    unit_key: str = "artist_id",
    track_key: str = "track_id",
    target: float = HOLDOUT_TARGET,
) -> SplitResult:
    """PREREG §3.1(3) candidate S3, applied to the frozen manifest rows.

    ``rows`` is the frozen manifest: any sequence of mappings carrying
    ``artist_id`` and ``track_id``.  The order of ``rows`` does not matter — the
    rule sorts everything it touches — so a manifest re-ordered on disk produces
    a byte-identical split.
    """
    groups = _index_by_unit(rows, unit_key, track_key)
    n_total = sum(len(t) for t in groups.values())
    threshold = target * n_total

    # smallest catalogue first, tie broken by unit id ascending
    order = sorted(groups, key=lambda a: (len(groups[a]), a))

    hold: list[str] = []
    cnt = 0
    walk: list[dict] = []
    for unit in order:
        if cnt >= threshold:  # STOP AT THE TARGET
            break
        hold.append(unit)
        cnt += len(groups[unit])
        walk.append(
            {"unit": unit, "tracks": len(groups[unit]), "cumulative": cnt}
        )

    held_units = tuple(sorted(hold))
    train_units = tuple(sorted(set(groups) - set(hold)))

    holdout_tracks = tuple(t for u in held_units for t in groups[u])
    training_tracks = tuple(t for u in train_units for t in groups[u])

    return SplitResult(
        rule="S3",
        rule_verbatim=RULE_VERBATIM_S3,
        unit_kind="artist",
        target_fraction=target,
        n_total=n_total,
        holdout_track_ids=tuple(sorted(holdout_tracks)),
        training_track_ids=tuple(sorted(training_tracks)),
        holdout_units=held_units,
        training_units=train_units,
        walk=tuple(walk),
    )


def by_track_split(
    rows: Sequence[Mapping],
    *,
    track_key: str = "track_id",
    target: float = HOLDOUT_TARGET,
    seed: int = BY_TRACK_SEED,
) -> SplitResult:
    """PREREG §3.3 — the single-artist (arm 3) intra-artist split.

    Registered exactly: ``numpy.random.default_rng(42).permutation`` over track
    ids **sorted ascending**, taking tracks until the count first reaches
    ``ceil(0.20 * N)``.  Arm 1 never calls this, which is why numpy is imported
    here rather than at module scope.
    """
    import numpy as np

    tracks = sorted(row[track_key] for row in rows)
    n_total = len(tracks)
    want = math.ceil(target * n_total)

    perm = np.random.default_rng(seed).permutation(n_total)
    holdout = [tracks[int(i)] for i in perm[:want]]
    held = set(holdout)
    training = [t for t in tracks if t not in held]

    return SplitResult(
        rule="BY_TRACK",
        rule_verbatim=RULE_VERBATIM_BY_TRACK,
        unit_kind="track",
        target_fraction=target,
        n_total=n_total,
        holdout_track_ids=tuple(sorted(holdout)),
        training_track_ids=tuple(training),
        holdout_units=tuple(sorted(holdout)),
        training_units=tuple(training),
    )


def pinned_stack() -> dict[str, str]:
    """The instrument behind ``by_track_split``'s permutation, for the freeze.

    ``numpy`` reads ``not installed on this box`` when the module is absent.
    That is the honest value for arm 1's freeze — the S3 walk is deterministic
    and RNG-free, so nothing it produced could have come from a random number —
    and it is a STOP for arm 3, whose split cannot run without it.
    """
    import platform

    try:
        import numpy as np

        numpy_version = np.__version__
    except ImportError:
        numpy_version = "not installed on this box"

    return {
        "python": platform.python_version(),
        "numpy": numpy_version,
        "by_track_seed": str(BY_TRACK_SEED),
        "note": (
            "numpy is consumed ONLY by by_track_split (arm 3). "
            "s3_artist_split is deterministic and RNG-free, so an arm-1 freeze "
            "on a box with no numpy is correct, not degraded."
        ),
    }