Ben Stull
99 lines
3.0 KiB
Python
99 lines
3.0 KiB
Python
"""Nearest-match selection of a catalog record for a knob coordinate."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from dataclasses import dataclass
|
|
|
|
from hef.catalog import Record
|
|
|
|
SELECTOR_MODES = frozenset({"none", "audio", "video", "av"})
|
|
CONTENT_MODES = frozenset({"audio", "video", "av"})
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Coordinate:
|
|
left: int
|
|
right: int
|
|
dark: int
|
|
light: int
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Weights:
|
|
brain: float = 1.0
|
|
mood: float = 1.0
|
|
|
|
|
|
def distance(a: Coordinate, b: Coordinate, weights: Weights = Weights()) -> float:
|
|
"""Weighted Euclidean distance; brain plane and mood plane weighted separately."""
|
|
brain_sq = (a.left - b.left) ** 2 + (a.right - b.right) ** 2
|
|
mood_sq = (a.dark - b.dark) ** 2 + (a.light - b.light) ** 2
|
|
return math.sqrt(weights.brain * brain_sq + weights.mood * mood_sq)
|
|
|
|
|
|
def record_coordinate(record: Record) -> Coordinate:
|
|
"""The (left, right, dark, light) coordinate of a catalog record."""
|
|
return Coordinate(record.left, record.right, record.dark, record.light)
|
|
|
|
|
|
def candidates_for_mode(records, mode: str, pool_size: int) -> list[Record]:
|
|
"""Records eligible for a content mode.
|
|
|
|
For 'av', if fewer than pool_size native 'av' records exist, fall back to
|
|
including 'audio' and 'video' records so the pool is never starved.
|
|
"""
|
|
if mode not in CONTENT_MODES:
|
|
raise ValueError(
|
|
f"candidates_for_mode expects a content mode {sorted(CONTENT_MODES)}, "
|
|
f"got {mode!r}"
|
|
)
|
|
primary = [r for r in records if r.mode == mode]
|
|
if mode == "av" and len(primary) < pool_size:
|
|
extra = [r for r in records if r.mode in {"audio", "video"}]
|
|
return primary + extra
|
|
return primary
|
|
|
|
|
|
from typing import Optional
|
|
|
|
|
|
def select(
|
|
records,
|
|
coord: Coordinate,
|
|
mode: str,
|
|
*,
|
|
pool_size: int = 4,
|
|
weights: Weights = Weights(),
|
|
approved_only: bool = False,
|
|
rng=None,
|
|
) -> Optional[Record]:
|
|
"""Pick the record nearest `coord` for the given selector `mode`.
|
|
|
|
- 'none' selector mode returns None (the void/rest state).
|
|
- With rng=None, returns the single nearest record (ties broken by id) for
|
|
deterministic behavior. Pass a random.Random to shuffle within the
|
|
pool_size nearest records.
|
|
- approved_only restricts to records the human has blessed.
|
|
"""
|
|
if mode not in SELECTOR_MODES:
|
|
raise ValueError(
|
|
f"invalid selector mode {mode!r}; expected one of {sorted(SELECTOR_MODES)}"
|
|
)
|
|
if mode == "none":
|
|
return None
|
|
pool = records
|
|
if approved_only:
|
|
pool = [r for r in pool if r.review_status == "approved"]
|
|
candidates = candidates_for_mode(pool, mode, pool_size)
|
|
if not candidates:
|
|
return None
|
|
ranked = sorted(
|
|
candidates,
|
|
key=lambda r: (distance(coord, record_coordinate(r), weights), r.id),
|
|
)
|
|
nearest = ranked[:pool_size]
|
|
if rng is None:
|
|
return nearest[0]
|
|
return rng.choice(nearest)
|