From 63e292a995305c688c6b667b493cc5fff781e270 Mon Sep 17 00:00:00 2001
From: Peter Woolery <peter@research.bike>
Date: Wed, 6 May 2026 06:48:26 -0700
Subject: [PATCH] feat: planner top-K set enumeration with protein/cuisine
 diversity penalty

---
 backend/app/services/planner/select.py | 53 +++++++++++++++++
 backend/tests/test_planner_select.py   | 78 ++++++++++++++++++++++++++
 2 files changed, 131 insertions(+)
 create mode 100644 backend/app/services/planner/select.py
 create mode 100644 backend/tests/test_planner_select.py

diff --git a/backend/app/services/planner/select.py b/backend/app/services/planner/select.py
new file mode 100644
index 0000000..228dc63
--- /dev/null
+++ b/backend/app/services/planner/select.py
@@ -0,0 +1,53 @@
+"""Top-K set enumeration with diversity penalty.
+
+Takes the top K (=20) scored recipes, enumerates all C(K, set_size)
+combinations, applies a pairwise diversity penalty for shared
+protein and cuisine, and returns the highest-scoring combination.
+"""
+from __future__ import annotations
+
+from itertools import combinations
+from typing import Iterable, List, Tuple
+
+from app.services.planner.config import PlannerConfig
+from app.services.planner.types import ScoredRecipe
+
+
+def set_diversity_penalty(
+    chosen: List[ScoredRecipe],
+    config: PlannerConfig,
+) -> float:
+    penalty = 0.0
+    for a, b in combinations(chosen, 2):
+        if a.protein and b.protein and a.protein == b.protein:
+            penalty += config.p_protein
+        if a.cuisine and b.cuisine and a.cuisine == b.cuisine:
+            penalty += config.p_cuisine
+    return penalty
+
+
+def _set_score(chosen: List[ScoredRecipe], config: PlannerConfig) -> float:
+    return sum(s.score for s in chosen) - set_diversity_penalty(chosen, config)
+
+
+def select_set(
+    scored: Iterable[ScoredRecipe],
+    config: PlannerConfig,
+) -> Tuple[List[ScoredRecipe], float]:
+    pool = list(scored)
+    if not pool:
+        return [], 0.0
+    if len(pool) <= config.set_size:
+        return pool, _set_score(pool, config)
+
+    pool.sort(key=lambda s: s.score, reverse=True)
+    candidate_pool = pool[: config.top_k]
+
+    best: List[ScoredRecipe] = []
+    best_score = float("-inf")
+    for combo in combinations(candidate_pool, config.set_size):
+        s = _set_score(list(combo), config)
+        if s > best_score:
+            best_score = s
+            best = list(combo)
+    return best, best_score
diff --git a/backend/tests/test_planner_select.py b/backend/tests/test_planner_select.py
new file mode 100644
index 0000000..baa1834
--- /dev/null
+++ b/backend/tests/test_planner_select.py
@@ -0,0 +1,78 @@
+from decimal import Decimal
+from uuid import uuid4
+
+from app.services.planner.config import PlannerConfig
+from app.services.planner.select import select_set, set_diversity_penalty
+from app.services.planner.types import RecipeCost, ScoredRecipe
+
+
+_CFG = PlannerConfig()
+
+
+def _mk(score, protein, cuisine):
+    rid = uuid4()
+    return ScoredRecipe(
+        recipe_id=rid,
+        score=score,
+        components={},
+        cost=RecipeCost(
+            recipe_id=rid,
+            total_cost=Decimal("10"),
+            total_savings=Decimal("1"),
+            sale_ingredient_count=1,
+            matched_ingredient_count=1,
+            total_ingredient_count=1,
+            pantry_hit_count=0,
+            line_items=[],
+        ),
+        protein=protein,
+        cuisine=cuisine,
+    )
+
+
+def test_diversity_penalty_zero_when_all_unique():
+    a = _mk(0.5, "chicken", "american")
+    b = _mk(0.5, "beef", "mexican")
+    c = _mk(0.5, "fish", "italian")
+    assert set_diversity_penalty([a, b, c], _CFG) == 0.0
+
+
+def test_diversity_penalty_three_chickens():
+    a = _mk(0.5, "chicken", "american")
+    b = _mk(0.5, "chicken", "italian")
+    c = _mk(0.5, "chicken", "mexican")
+    # 3 protein pairs * 0.15 = 0.45, no cuisine pairs
+    p = set_diversity_penalty([a, b, c], _CFG)
+    assert abs(p - 0.45) < 1e-6
+
+
+def test_select_set_picks_diverse_over_homogeneous():
+    """Three chicken/american (each 0.95) lose to fully-diverse mixed
+    (each 0.90) once protein and cuisine penalties apply.
+
+    Mix candidates each overlap high on exactly one attribute so that
+    *any* "1 high + 2 mix" pairing also incurs penalty -- otherwise the
+    optimum would mix a single high with two unrelated mixes."""
+    high1 = _mk(0.95, "chicken", "american")
+    high2 = _mk(0.95, "chicken", "american")
+    high3 = _mk(0.95, "chicken", "american")
+    mix1 = _mk(0.90, "pork", "thai")          # fully distinct from high
+    mix2 = _mk(0.90, "chicken", "mexican")    # shares protein with high
+    mix3 = _mk(0.90, "fish", "american")      # shares cuisine with high
+
+    chosen, set_score = select_set([high1, high2, high3, mix1, mix2, mix3], _CFG)
+    chosen_ids = {s.recipe_id for s in chosen}
+    assert chosen_ids == {mix1.recipe_id, mix2.recipe_id, mix3.recipe_id}
+
+
+def test_select_set_handles_too_few():
+    a = _mk(0.5, "x", "y")
+    b = _mk(0.5, "x", "y")
+    chosen, _ = select_set([a, b], _CFG)
+    assert len(chosen) == 2  # less than set_size returns what we have
+
+
+def test_select_set_empty_input():
+    chosen, score = select_set([], _CFG)
+    assert chosen == []
+    assert score == 0.0