Extract cast context policy

cast_context.py

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+from typing import Any, Callable
+
+
+Choose = Callable[[Any, list[tuple[str, str, str]]], tuple[str, str, str]]
+
+
+def count_phrase(count: int, singular: str, plural: str) -> str:
+    words = {
+        0: "no",
+        1: "one",
+        2: "two",
+        3: "three",
+        4: "four",
+        5: "five",
+        6: "six",
+        7: "seven",
+        8: "eight",
+        9: "nine",
+        10: "ten",
+        11: "eleven",
+        12: "twelve",
+    }
+    label = singular if count == 1 else plural
+    return f"{words.get(count, str(count))} {label}"
+
+
+def cast_summary_phrase(women_count: int, men_count: int) -> str:
+    women_count = max(0, int(women_count))
+    men_count = max(0, int(men_count))
+    if women_count + men_count == 0:
+        women_count = 1
+    person_count = women_count + men_count
+    women_label = "woman" if women_count == 1 else "women"
+    men_label = "man" if men_count == 1 else "men"
+    return f"{women_count} {women_label}, {men_count} {men_label}, {person_count} total adults"
+
+
+def configured_cast_context(women_count: int, men_count: int) -> dict[str, str]:
+    women_count = max(0, int(women_count))
+    men_count = max(0, int(men_count))
+    if women_count + men_count == 0:
+        women_count = 1
+    parts = []
+    if women_count:
+        parts.append(count_phrase(women_count, "adult woman", "adult women"))
+    if men_count:
+        parts.append(count_phrase(men_count, "adult man", "adult men"))
+    subject_phrase = parts[0] if len(parts) == 1 else f"{parts[0]} and {parts[1]}"
+    person_count = women_count + men_count
+    if person_count == 1:
+        scene_kind = "solo adult sexual pose"
+    elif person_count == 2:
+        scene_kind = "adult couple sex scene"
+    elif person_count == 3:
+        scene_kind = "adult threesome sex scene"
+    else:
+        scene_kind = "adult group sex scene"
+    return {
+        "subject_type": "configured_cast",
+        "subject": f"{women_count}w_{men_count}m_sex_scene",
+        "subject_phrase": subject_phrase,
+        "age": "21+ adults",
+        "body": "varied",
+        "skin": "",
+        "hair": "",
+        "eyes": "",
+        "body_phrase": "varied adult bodies",
+        "women_count": str(women_count),
+        "men_count": str(men_count),
+        "person_count": str(person_count),
+        "cast_summary": cast_summary_phrase(women_count, men_count),
+        "scene_kind": scene_kind,
+    }
+
+
+def couple_type_from_counts(
+    rng: Any,
+    women_count: int,
+    men_count: int,
+    *,
+    choose: Choose,
+    couple_types: list[tuple[str, str, str]],
+) -> tuple[str, str, str, int, int]:
+    women_count = max(0, int(women_count))
+    men_count = max(0, int(men_count))
+    if women_count >= 2 and men_count == 0:
+        return "two women", "two women", "close affectionate couple pose", 2, 0
+    if men_count >= 2 and women_count == 0:
+        return "two men", "two men", "relaxed romantic couple pose", 0, 2
+    if women_count >= 1 and men_count >= 1:
+        return "woman and man", "a woman and a man", "playful date-night pose", 1, 1
+
+    primary_subject, subject_phrase, pose = choose(rng, couple_types)
+    if primary_subject == "two women":
+        return primary_subject, subject_phrase, pose, 2, 0
+    if primary_subject == "two men":
+        return primary_subject, subject_phrase, pose, 0, 2
+    return primary_subject, subject_phrase, pose, 1, 1