"""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