recommender.py

# src/tft_analyzer/recommender.py
from typing import Dict, List, Any


def get_trait_breakpoints(trait_name: str, trait_data: Dict[str, Any]) -> List[int]:
    """Returns a sorted list of minUnits for a trait."""
    if trait_name not in trait_data:
        return []
    effects = trait_data[trait_name].get("effects", [])
    if not effects:
        return []
    return sorted([int(e.get("minUnits", 0)) for e in effects])


def evaluate_comp(
    comp: List[str],
    champion_data: Dict[str, Dict[str, Any]],
    trait_data: Dict[str, Any],
    units_gone: Dict[str, Dict[str, int]],
    level_odds: Dict[int, float]
) -> float:
    """
    Calculates a heuristic score for a composition based on:
    1. Active Trait Breakpoints (High Value)
    2. Proximity to next Breakpoint (Medium Value)
    3. Champion Individual Value (Cost, Odds, Contention)
    """
    if not comp:
        return 0.0

    # --- 1. Trait Score ---
    trait_counts = {}
    for name in comp:
        if name not in champion_data:
            continue
        for trait in champion_data[name].get("traits", []):
            trait_counts[trait] = trait_counts.get(trait, 0) + 1

    trait_score = 0.0
    for trait, count in trait_counts.items():
        breakpoints = get_trait_breakpoints(trait, trait_data)
        if not breakpoints:
            continue

        # Find active breakpoint and next target
        active_min = 0
        next_min = breakpoints[0]

        for idx, bp in enumerate(breakpoints):
            if count >= bp:
                active_min = bp
                if idx + 1 < len(breakpoints):
                    next_min = breakpoints[idx+1]
                else:
                    next_min = None  # Already at max breakpoint
            else:
                next_min = bp
                break

        # Scoring Weights
        # Active traits are worth a lot (synergy power)
        trait_score += active_min * 15.0

        # Proximity score: encourage the algorithm to "finish" a trait
        if next_min is not None and next_min > 0 and count < next_min:
            trait_score += (count / next_min) * 5.0

    # --- 2. Champion Value (EV) ---
    champ_score = 0.0
    for name in comp:
        if name not in champion_data:
            continue

        cost = champion_data[name].get("cost", 1)
        tier_str = f"{cost}-cost"

        # Calculate contention penalty
        gone = units_gone.get(tier_str, {}).get(name, 0)

        # Probability of finding this unit tier
        prob = level_odds.get(cost, 0)

        # Heuristic Value Formula: (Cost^1.5 * Odds) / (1 + Contention)
        # We value high-cost units more, but only if odds are decent.
        val = (pow(cost, 1.5) * prob) / (1.0 + (gone * 0.2))

        champ_score += val * 10.0

    return trait_score + champ_score


def _jaccard(a: List[str], b: List[str]) -> float:
    """Jaccard similarity between two champion lists."""
    sa, sb = set(a), set(b)
    if not sa and not sb:
        return 1.0
    return len(sa & sb) / len(sa | sb)


def _get_trait_pairs(
    champion_data: Dict[str, Dict[str, Any]],
    available: List[str],
    locked: List[str],
) -> List[List[str]]:
    """
    Build seed pairs: for each trait with 3+ available champs,
    pick the two highest-cost available champs sharing that trait.
    Returns list of [champ1, champ2] pairs (excluding locked champs).
    """
    # Group available champs by trait
    trait_champs: Dict[str, List[str]] = {}
    for name in available:
        if name in locked or name not in champion_data:
            continue
        for trait in champion_data[name].get("traits", []):
            trait_champs.setdefault(trait, []).append(name)

    pairs = []
    for trait, champs in trait_champs.items():
        if len(champs) < 2:
            continue
        # Sort by cost descending to pick the strongest pair
        champs.sort(key=lambda n: champion_data[n].get("cost", 1), reverse=True)
        pairs.append(champs[:2])

    return pairs


def _run_greedy(
    seed: List[str],
    available: List[str],
    champion_data: Dict[str, Dict[str, Any]],
    trait_data: Dict[str, Any],
    units_gone: Dict[str, Dict[str, int]],
    level_odds: Dict[int, float],
    level: int,
) -> List[str]:
    """Run greedy fill from a seed comp up to the given level."""
    working = list(seed)
    while len(working) < level:
        best_candidate = None
        best_candidate_score = -float('inf')

        for name in available:
            if name in working:
                continue
            temp_comp = working + [name]
            score = evaluate_comp(temp_comp, champion_data, trait_data, units_gone, level_odds)
            if score > best_candidate_score:
                best_candidate_score = score
                best_candidate = name

        if best_candidate:
            working.append(best_candidate)
        else:
            break

    return working


def recommend_comp(
    champions_by_tier: Dict[str, List[str]],
    champion_data: Dict[str, Dict[str, Any]],
    trait_data: Dict[str, Any],
    level: int,
    units_gone: Dict[str, Dict[str, int]],
    level_odds: Dict[int, float],
    locked_champions: List[str] = None,
    top_n: int = 1,
) -> List[Dict[str, Any]]:
    """
    Recommends team compositions using a multi-seed greedy algorithm.

    Returns a list of scored comp dicts sorted by score descending:
        [{"champions": [...], "score": float, "source": "greedy"}, ...]

    When top_n=1, returns a single-element list (backward compatible shape).
    """
    if locked_champions is None:
        locked_champions = []

    # 1. Filter available champions based on current shop odds
    available = []
    for tier_str, names in champions_by_tier.items():
        tier_cost = int(tier_str.split('-')[0])
        if level_odds.get(tier_cost, 0) > 0:
            available.extend(names)

    for name in locked_champions:
        if name not in available:
            available.append(name)

    if not available:
        return []

    # 2. Build seeds
    scored_champs = []
    for name in available:
        if name not in champion_data:
            continue
        if name in locked_champions:
            continue
        cost = champion_data[name].get("cost", 1)
        prob = level_odds.get(cost, 0)
        score = cost * prob
        scored_champs.append((score, name))

    scored_champs.sort(key=lambda x: x[0], reverse=True)

    base = list(locked_champions)
    seeds: List[List[str]] = [list(base)]  # pure greedy seed

    # Top-carry seeds
    if len(base) < level:
        for _, name in scored_champs[:5]:
            if name not in base:
                seeds.append(list(base) + [name])

    # Trait-pair seeds (for diversity)
    if top_n > 1:
        for pair in _get_trait_pairs(champion_data, available, locked_champions):
            seed = list(base) + [n for n in pair if n not in base]
            if len(seed) <= level:
                seeds.append(seed)

    # 3. Run greedy for each seed, collect all results
    all_comps: List[Dict[str, Any]] = []
    for seed in seeds:
        comp = _run_greedy(seed, available, champion_data, trait_data, units_gone, level_odds, level)
        raw_score = evaluate_comp(comp, champion_data, trait_data, units_gone, level_odds)
        all_comps.append({"champions": comp, "raw_score": raw_score})

    if not all_comps:
        return []

    # 4. Normalize scores to 0-100
    max_score = max(c["raw_score"] for c in all_comps)
    if max_score <= 0:
        max_score = 1.0

    for c in all_comps:
        c["score"] = round((c["raw_score"] / max_score) * 100, 1)
        c["source"] = "greedy"
        del c["raw_score"]

    # 5. Sort by score descending
    all_comps.sort(key=lambda c: c["score"], reverse=True)

    # 6. Deduplicate (Jaccard > 0.8 = too similar)
    unique: List[Dict[str, Any]] = []
    for comp in all_comps:
        is_dup = False
        for existing in unique:
            if _jaccard(comp["champions"], existing["champions"]) > 0.8:
                is_dup = True
                break
        if not is_dup:
            unique.append(comp)
        if len(unique) >= top_n:
            break

    return unique