backend/assignment.py

"""Balanced random subset assignment: pure logic, no FastAPI/HF imports.

Everything here takes plain data in and returns plain data out, so it can be
unit-tested (see test_assignment.py) without touching the network. The I/O
layer (reading/writing annotation and assignment records from the HF dataset
repo, locking) lives in main.py and calls into these functions.

Config is env-overridable, matching the existing HF_* var convention in
main.py, so these three numbers can change without touching the algorithm.
"""

import os
import re
import secrets
import string
from datetime import datetime

TOTAL_VIDEOS = int(os.getenv("TOTAL_VIDEOS", "200"))
VIDEOS_PER_ANNOTATOR = int(os.getenv("VIDEOS_PER_ANNOTATOR", "20"))
TARGET_ANNOTATIONS_PER_VIDEO = int(os.getenv("TARGET_ANNOTATIONS_PER_VIDEO", "5"))
# How long an in-progress (not-yet-completed) round still counts as an open
# "reservation" against other annotators' balancing, before being treated as
# abandoned. Distinct from COMPLETION_GRACE_PERIOD_SECONDS below.
RESERVATION_TTL_SECONDS = int(os.getenv("ASSIGNMENT_RESERVATION_TTL_SECONDS", str(24 * 3600)))
# How long a just-completed round keeps showing its completion code before a
# later visit automatically gets the next round instead. Distinct from
# RESERVATION_TTL_SECONDS above - this one is "how long do you get to read
# your own code," not "when do other annotators stop waiting on your slot."
COMPLETION_GRACE_PERIOD_SECONDS = int(os.getenv("COMPLETION_GRACE_PERIOD_SECONDS", str(10 * 60)))

COMPLETION_CODE_ALPHABET = string.ascii_uppercase + string.digits
COMPLETION_CODE_LENGTH = 12


def generate_completion_code():
    """Cryptographically random completion code, format T2AV-XXXXXXXXXXXX.
    Uses `secrets`, not `random` - this is a credential the annotator submits
    elsewhere (e.g. an MTurk HIT) as proof of completion, not a UI cosmetic."""
    suffix = "".join(secrets.choice(COMPLETION_CODE_ALPHABET) for _ in range(COMPLETION_CODE_LENGTH))
    return f"T2AV-{suffix}"


def slugify(value):
    """Canonical annotator_id from a raw display name. Lowercased (unlike
    main.py's filename slugify) so "Adi" and "adi" resolve to one identity -
    this is the identity key used for completion/assignment lookups, not a
    filename, so case-insensitivity is the more correct default here."""
    slug = re.sub(r"[^A-Za-z0-9_.-]+", "-", value or "").strip("-").lower()
    return slug or "unknown"


def parse_iso(value):
    return datetime.fromisoformat(value.replace("Z", "+00:00"))


def dedupe_latest(annotation_records):
    """A "completed annotation" is one row per unique (annotator_id, video_id)
    pair - repeated saves/autosaves/updates from the same annotator for the
    same video collapse into a single record, keeping the most recent
    created_at. Records missing a required field are skipped (never crash on
    older/malformed data) and reported separately.

    Returns (deduped: dict[(annotator_id, video_id)] -> record, skipped: int).
    """
    latest = {}
    skipped = 0

    for record in annotation_records:
        try:
            video_id = record["video_id"]
            user_raw = record["user"]
            created_at = record["created_at"]
            if not video_id or not user_raw or not created_at:
                raise ValueError("empty required field")
            parse_iso(created_at)  # validate shape before trusting string comparisons below
        except (KeyError, TypeError, ValueError, AttributeError):
            skipped += 1
            continue

        key = (slugify(user_raw), video_id)
        existing = latest.get(key)
        if existing is None or created_at > existing["created_at"]:
            latest[key] = record

    return latest, skipped


def completed_by_video(deduped):
    result = {}
    for annotator_id, video_id in deduped:
        result.setdefault(video_id, set()).add(annotator_id)
    return result


def completed_by_annotator(deduped):
    result = {}
    for annotator_id, video_id in deduped:
        result.setdefault(annotator_id, set()).add(video_id)
    return result


def saved_responses_for_round(deduped, annotator_id, video_ids):
    """{video_id: responses} for every video in `video_ids` that `annotator_id`
    already has a saved annotation for - lets a returning annotator (same
    name, new session or a reload) resume with prior answers pre-filled
    instead of starting blank. `responses` is exactly the {events, video}
    shape the frontend's annotation state already uses, since it's what was
    saved from that same shape originally."""
    video_id_set = set(video_ids)
    return {
        video_id: record["annotations"]["responses"]
        for (a_id, video_id), record in deduped.items()
        if a_id == annotator_id and video_id in video_id_set
    }


def reserved_by_video(assignment_records, completed_by_annotator_map, now, ttl_seconds, exclude_annotator=None):
    """Open (not-yet-completed, not-yet-expired) reservations per video,
    across every OTHER annotator's assignment. A reservation stops counting
    once its assignment is older than ttl_seconds ("abandoned") - this never
    deletes or mutates the stored assignment record, it only stops
    contributing to the live coverage snapshot used for balancing."""
    result = {}
    for record in assignment_records:
        annotator_id = record.get("annotator_id")
        if not annotator_id or annotator_id == exclude_annotator:
            continue
        try:
            created_at = parse_iso(record["created_at"])
        except (KeyError, ValueError, AttributeError, TypeError):
            continue
        if (now - created_at).total_seconds() > ttl_seconds:
            continue

        completed = completed_by_annotator_map.get(annotator_id, set())
        for video_id in record.get("video_ids", []):
            if video_id in completed:
                continue
            result[video_id] = result.get(video_id, 0) + 1
    return result


def pick_balanced(candidates, coverage_fn, count, rng, target=TARGET_ANNOTATIONS_PER_VIDEO):
    """Picks `count` videos, prioritizing lowest coverage (completed +
    reserved) first, random tie-break among equal coverage. Prefers
    strictly-under-target videos; only reaches into at-or-over-target ones if
    there aren't enough under-target candidates to satisfy `count`."""
    if count <= 0:
        return []

    under_target = [v for v in candidates if coverage_fn(v) < target]
    pool = under_target if len(under_target) >= count else list(candidates)

    ordered = sorted(pool, key=lambda video_id: (coverage_fn(video_id), rng.random()))
    return ordered[:count]


def build_round(
    annotator_id,
    annotator_raw,
    catalog_video_ids,
    already_completed,
    completed_by_video_map,
    reserved_by_video_map,
    rng,
    round_number,
    videos_per_annotator=VIDEOS_PER_ANNOTATOR,
    target=TARGET_ANNOTATIONS_PER_VIDEO,
):
    """Builds a new round: exactly `videos_per_annotator` videos this
    annotator has never completed before - across every prior round AND any
    completions that predate this feature (`already_completed` is the full,
    round-agnostic history). Videos already completed are unconditionally
    excluded, never used to seed the round. Returns fewer than
    `videos_per_annotator` only when the eligible pool itself is smaller
    (see `actual_size` vs `requested_size` in the result) - never fills the
    gap with already-completed videos."""
    candidates = [v for v in catalog_video_ids if v not in already_completed]

    def coverage(video_id):
        return len(completed_by_video_map.get(video_id, set())) + reserved_by_video_map.get(video_id, 0)

    video_ids = pick_balanced(candidates, coverage, videos_per_annotator, rng, target)

    return {
        "annotator_id": annotator_id,
        "annotator_name_raw": annotator_raw,
        "round_number": round_number,
        "video_ids": video_ids,
        "requested_size": videos_per_annotator,
        "actual_size": len(video_ids),
    }


def summarize(
    catalog_video_ids,
    completed_by_video_map,
    reserved_by_video_map,
    total_completed_preserved,
    skipped_records,
    duplicate_records_deduped,
    assignment_count,
    target=TARGET_ANNOTATIONS_PER_VIDEO,
):
    below = sum(1 for v in catalog_video_ids if len(completed_by_video_map.get(v, set())) < target)
    return {
        "total_videos": len(catalog_video_ids),
        "target_per_video": target,
        "videos_per_annotator": VIDEOS_PER_ANNOTATOR,
        "completed_counts": {v: len(completed_by_video_map.get(v, set())) for v in catalog_video_ids},
        "reserved_counts": {v: reserved_by_video_map.get(v, 0) for v in catalog_video_ids},
        "videos_below_target": below,
        "videos_at_or_above_target": len(catalog_video_ids) - below,
        "total_existing_annotations_preserved": total_completed_preserved,
        "skipped_malformed_records": skipped_records,
        "duplicate_records_deduplicated": duplicate_records_deduped,
        "total_assignments": assignment_count,
    }