Blame for server/deckbuilder.py

"""
deckbuilder.py

Turns "commander(s) + partial decklist" into a complete 100-card Commander
deck using the production ensemble (training/ensemble.py). The model only
ever proposes nonbasic cards -- basic lands are excluded from the tokenizer
vocab -- so this module owns everything around it:

1. Resolve pasted names against the vocab (exact, then fuzzy with a
   warning), keeping user-supplied basics aside since the vocab can't
   represent them.
2. Ask the ensemble for every open slot, with Commander-banned cards
   excluded from its candidate pool.
3. Hit a land target: trim the ensemble's lowest-scoring nonland,
   non-staple picks and replace them with basics split by the deck's
   colored-pip distribution.

Land target: DEFAULT_TARGET_LANDS (38) unless the caller sets one, and never
below MIN_TARGET_LANDS (36) -- the project owner's deckbuilding standard.
(The training data suggests players run fewer: estimated median totals of
33-35 by color count. That dataset drops basic lands from ~92% of its decks,
10,300 of 122,571 contain any basic, which is also why the ensemble fills
open slots almost entirely with nonbasics and needs this step.) The target
is a minimum: lands the model picks beyond it are kept.
"""

from __future__ import annotations

import secrets
import threading
from collections import Counter
from dataclasses import dataclass, field
from difflib import get_close_matches

import torch

from tokenizer.mtg_tokenizer import BASIC_LAND_NAMES, COLORS, CardRecord, CardTokenizer, _normalize_name, _pip_counts
from server.card_prices import PriceBook
from server.decklist import parse_decklist

DECK_SIZE = 100
DEFAULT_TARGET_LANDS = 38
# Cards picked per model pass in the synergy phase and prune refills: an int
# for a fixed size, or "dynamic" to vary it each pass with how alike the
# model's top candidates are (see training/complete_deck.py's
# _choose_dynamic_chunk_size). Set here rather than inherited from
# ensemble_complete_deck's default, so the live behavior only changes when
# this line does.
GENERATION_CHUNK_SIZE: int | str = "dynamic"
MIN_TARGET_LANDS = 36
BASIC_FOR_COLOR = {"W": "Plains", "U": "Island", "B": "Swamp", "R": "Mountain", "G": "Forest"}
BASIC_TYPE_LINES = {name: f"Basic Land — {name}" for name in BASIC_LAND_NAMES - {"Wastes"}} | {"Wastes": "Basic Land"}
BASIC_COLOR = {v: k for k, v in BASIC_FOR_COLOR.items()}
BASIC_TEXT = {name: f"({{T}}: Add {{{BASIC_COLOR.get(name, 'C')}}}.)" for name in BASIC_LAND_NAMES}

# Grouping buckets; a card goes in the first one its FRONT face's type line
# matches. Land comes first so artifact lands and Dryad Arbor are grouped
# (and counted in the summary) the same way the land-target math counts them.
CATEGORIES = ["Land", "Creature", "Planeswalker", "Battle", "Instant", "Sorcery", "Artifact", "Enchantment"]

class DeckInputError(ValueError):
    """Bad user input that the API should report back (HTTP 422), not crash
    on. `unresolved` lists names that matched nothing, each with
    suggestions, so the UI can offer one-click fixes."""

    def __init__(self, message: str, unresolved: list[dict] | None = None):
        super().__init__(message)
        self.message = message
        self.unresolved = unresolved or []


@dataclass
class DeckCard:
    name: str
    quantity: int
    source: str  # "input" | "synergy" | "staple" | "basic"
    score: float | None
    type_line: str
    mana_cost: str
    cmc: float
    category: str
    oracle_text: str
    locked: bool = False  # kept from a previous completion at the user's request
    # Real-deck-frequency classification ("staple" | "mid" | "long_tail"),
    # None for basics (outside the tokenizer vocab, so untiered). Distinct
    # from `source`: the synergy phase runs unrestricted and routinely
    # picks real staples (Sol Ring, Command Tower) on their own merit, so
    # a card can be source="synergy" and tier="staple" at once -- without
    # this field the frontend had no way to show that, and a staple could
    # read as "just a synergy pick" with nothing marking it as universal.
    tier: str | None = None
    # Per-face power/toughness, loyalty, or defense ("4/3", "Loyalty 3", None), see server/card_stats.py.
    stats: list[str | None] = field(default_factory=list)
    # Cheapest paper printing: {"usd": 0.25, "set": "M21", "finish": "foil"}, see server/card_prices.py.
    price: dict | None = None
    color_identity: list[str] = field(default_factory=list)
    # Whether a deck may hold more than one copy (basics, Relentless Rats, ...).
    multiples_allowed: bool = False

    def to_dict(self) -> dict:
        return self.__dict__.copy()


@dataclass
class Completion:
    commanders: list[DeckCard]
    cards: list[DeckCard]
    warnings: list[str] = field(default_factory=list)
    target_lands: int = 0
    prices_updated: str | None = None

    def to_dict(self) -> dict:
        land_count = sum(c.quantity for c in self.cards if c.category == "Land")
        everything = self.commanders + self.cards
        return {
            "commanders": [c.to_dict() for c in self.commanders],
            "cards": [c.to_dict() for c in self.cards],
            "warnings": self.warnings,
            "prices_updated": self.prices_updated,
            "summary": {
                "price_usd": round(sum(c.price["usd"] * c.quantity for c in everything if c.price), 2),
                "unpriced": sum(c.quantity for c in everything if not c.price),
                "total": sum(c.quantity for c in self.commanders + self.cards),
                "lands": land_count,
                "target_lands": self.target_lands,
                "added": sum(c.quantity for c in self.cards if c.source != "input"),
                "basics_added": sum(c.quantity for c in self.cards if c.source == "basic"),
            },
        }


def front_face_type(card: CardRecord) -> str:
    return card.type_line.split("//")[0]


def category_of(type_line: str) -> str:
    front = type_line.split("//")[0]
    for cat in CATEGORIES:
        if cat in front:
            return cat
    return "Other"


def is_front_land(card: CardRecord) -> bool:
    return "Land" in front_face_type(card)


def is_commander_eligible(card: CardRecord) -> bool:
    front = front_face_type(card)
    return (
        ("Legendary" in front and ("Creature" in front or "Background" in front))
        or "can be your commander" in card.oracle_text
    )


def _allows_multiple(card: CardRecord) -> bool:
    text = card.oracle_text
    return "Basic" in card.type_line or "A deck can have any number" in text or "A deck can have up to" in text


def _partner_pair_ok(a: CardRecord, b: CardRecord) -> bool:
    def has(card: CardRecord, phrase: str) -> bool:
        return phrase in card.oracle_text or phrase in card.keywords

    if has(a, "Partner with") or has(b, "Partner with"):
        return f"Partner with {b.name}" in a.oracle_text and f"Partner with {a.name}" in b.oracle_text
    for phrase in ("Partner", "Friends forever"):
        if has(a, phrase) and has(b, phrase):
            return True
    for x, y in ((a, b), (b, a)):
        if has(x, "Choose a Background") and "Background" in front_face_type(y):
            return True
        if has(x, "Doctor's companion") and "Doctor" in front_face_type(y):
            return True
    return False


def allocate_basics(n: int, colors: set[str], pip_totals: dict[str, float]) -> Counter:
    """Split n basic lands across the commander's colors in proportion to
    the deck's colored pips (largest-remainder rounding). Colorless decks
    get Wastes. Colors with no pips at all still get a share when every
    color has zero pips (e.g. an all-colorless-spells deck)."""
    if n <= 0:
        return Counter()
    if not colors:
        return Counter({"Wastes": n})
    ordered = [c for c in COLORS if c in colors]
    weights = [pip_totals.get(c, 0.0) for c in ordered]
    if sum(weights) == 0:
        weights = [1.0] * len(ordered)
    total = sum(weights)
    exact = [n * w / total for w in weights]
    counts = [int(x) for x in exact]
    by_remainder = sorted(range(len(ordered)), key=lambda i: exact[i] - counts[i], reverse=True)
    for i in by_remainder[: n - sum(counts)]:
        counts[i] += 1
    return Counter({BASIC_FOR_COLOR[c]: k for c, k in zip(ordered, counts) if k > 0})


class DeckBuilder:
    """Loads the three checkpoints once and serves completions. Inference is
    serialized with a lock: requests arrive on FastAPI's thread pool, and a
    completion is only ~1-3s, so a queue of one is simpler than per-thread
    model copies.

    `version` is the Hub release record from server/model_files.py (None
    when the local files were never synced), reported by /api/health.
    `card_stats` is the preview lookup from server/card_stats.py; the app
    fills it in after loading, since building it needs the tokenizer.
    `prices` serves cheapest-printing prices (server/card_prices.py); each
    completion nudges it to refresh in the background when stale."""

    def __init__(
        self,
        version: dict | None = None,
        card_stats: dict[str, list[str | None]] | None = None,
        prices: PriceBook | None = None,
    ):
        from training.ensemble import PRUNE_CHECKPOINT, STAPLE_CHECKPOINT, SYNERGY_CHECKPOINT, load_tiers
        from training.evaluate import load_checkpoint
        from training.train import pick_device
        from training.train_prune import load_prune_checkpoint

        self.device = pick_device()
        self.staple_model, self.tok, _ = load_checkpoint(STAPLE_CHECKPOINT, self.device)
        self.synergy_model, _, _ = load_checkpoint(SYNERGY_CHECKPOINT, self.device)
        self.prune_model, _, _ = load_prune_checkpoint(PRUNE_CHECKPOINT, self.device)
        self.tiers = load_tiers(self.tok)
        # A release's tokenizer can list cards newer than its checkpoints
        # (new cards are appended to the end of the vocab), and ids past a
        # checkpoint's vocab must never reach it. The model only ever
        # suggests cards below this; newer ones can still be in a deck.
        models = [self.staple_model, self.synergy_model, self.prune_model]
        self.model_vocab_size = min(m.card_embedding.vocab_size for m in models)
        self.version = version
        self.card_stats = card_stats or {}
        self.prices = prices or PriceBook(allow_download=False, data={})
        self._lock = threading.Lock()
        self._search_index = self._build_search_index(self.tok)

    # ------------------------------------------------------------------
    # Card lookup
    # ------------------------------------------------------------------

    @staticmethod
    def _build_search_index(tok: CardTokenizer) -> list[tuple[str, CardRecord]]:
        index = [(_normalize_name(card.name), card) for card in tok.card_by_token.values()]
        index.sort(key=lambda pair: pair[0])
        return index

    def search(
        self, query: str, commander_only: bool = False, limit: int = 10,
        colors: str | None = None, basics: bool = False,
    ) -> list[dict]:
        """Autocomplete: whole-name prefix matches first, then word-prefix,
        then substring matches, each alphabetical.

        `colors` (e.g. "BG"; "" for colorless) keeps only cards whose color
        identity fits within it, for adding cards to a deck. `basics` also
        offers matching basic lands, which aren't in the tokenizer's vocab
        (first, since they're exact prefix matches)."""
        q = _normalize_name(query)
        if not q:
            return []
        allowed = set(colors) if colors is not None else None
        basic_matches = [
            self._basic_summary(name) for name in sorted(BASIC_LAND_NAMES)
            if basics and not commander_only and _normalize_name(name).startswith(q)
            and (allowed is None or set(self._basic_card(name, 1, "input").color_identity) <= allowed)
        ]
        buckets: tuple[list, list, list] = ([], [], [])
        for norm, card in self._search_index:
            if commander_only and not is_commander_eligible(card):
                continue
            if allowed is not None and not set(card.color_identity) <= allowed:
                continue
            if norm.startswith(q):
                buckets[0].append(card)
            elif f" {q}" in norm:
                buckets[1].append(card)
            elif q in norm:
                buckets[2].append(card)
        results = (buckets[0] + buckets[1] + buckets[2])[: max(0, limit - len(basic_matches))]
        return basic_matches + [self._card_summary(card) for card in results]

    def _basic_summary(self, name: str) -> dict:
        card = self._basic_card(name, 1, "input")
        return {
            "name": name, "type_line": card.type_line, "mana_cost": "", "cmc": 0.0, "category": "Land",
            "oracle_text": card.oracle_text, "stats": [], "price": card.price, "tier": None,
            "color_identity": card.color_identity, "multiples_allowed": True, "banned": False,
        }

    def _card_summary(self, card: CardRecord) -> dict:
        """A search result: enough to show the card (cost, rules text, stats,
        price) anywhere a deck card can be shown, e.g. in the ban list."""
        return {
            "name": card.name,
            "type_line": card.type_line,
            "mana_cost": card.mana_cost,
            "cmc": card.cmc,
            "category": category_of(card.type_line),
            "oracle_text": card.oracle_text,
            "stats": self.card_stats.get(card.oracle_id, []),
            "price": self.prices.get(card.oracle_id),
            "tier": self.tiers.get(self.tok.token_to_id[self.tok.oracle_id_to_token[card.oracle_id]]),
            "color_identity": card.color_identity,
            "multiples_allowed": _allows_multiple(card),
            "banned": card.commander_legality == "banned",
        }

    def _card(self, name: str) -> CardRecord:
        return self.tok.card_by_token[self.tok.resolve_name(name)]

    def _suggestions(self, name: str, commander_only: bool = False) -> list[str]:
        norm = _normalize_name(name)
        candidates = get_close_matches(norm, self.tok._name_index.keys(), n=8, cutoff=0.6)
        names: list[str] = []
        for c in candidates:
            card = self.tok.card_by_token[self.tok._name_index[c]]
            if commander_only and not is_commander_eligible(card):
                continue
            if card.name not in names:
                names.append(card.name)
        if not names:
            names = [r["name"] for r in self.search(name, commander_only=commander_only, limit=3)]
        return names[:3]

    def _resolve(self, name: str, warnings: list[str]) -> CardRecord | None:
        result = self.tok.resolve_name_with_confidence(name)
        if result is None:
            return None
        token, exact = result
        card = self.tok.card_by_token[token]
        if not exact:
            warnings.append(f'Read "{name}" as "{card.name}".')
        return card

    def _deck_card(self, card: CardRecord, quantity: int, source: str, score: float | None) -> DeckCard:
        return DeckCard(
            name=card.name, quantity=quantity, source=source, score=score,
            type_line=card.type_line, mana_cost=card.mana_cost, cmc=card.cmc,
            category=category_of(card.type_line), oracle_text=card.oracle_text,
            stats=self.card_stats.get(card.oracle_id, []), price=self.prices.get(card.oracle_id),
            tier=self.tiers.get(self.tok.token_to_id[self.tok.oracle_id_to_token[card.oracle_id]]),
            color_identity=card.color_identity, multiples_allowed=_allows_multiple(card),
        )

    def _basic_card(self, name: str, quantity: int, source: str) -> DeckCard:
        return DeckCard(
            name=name, quantity=quantity, source=source, score=None,
            type_line=BASIC_TYPE_LINES[name], mana_cost="", cmc=0.0, category="Land",
            oracle_text=BASIC_TEXT[name], price=self.prices.basic(name),
            color_identity=[BASIC_COLOR[name]] if name in BASIC_COLOR else [], multiples_allowed=True,
        )

    # ------------------------------------------------------------------
    # Completion
    # ------------------------------------------------------------------

    def model_knows(self, card: CardRecord) -> bool:
        return self.tok.token_to_id[self.tok.oracle_id_to_token[card.oracle_id]] < self.model_vocab_size

    def complete(
        self,
        commander_names: list[str],
        decklist: str,
        target_lands: int | None = None,
        prune: bool = True,
        locked: list[tuple[str, int, str, bool]] | None = None,
        exclude: list[str] | None = None,
    ) -> Completion:
        """`locked`: (name, quantity, source) cards kept from a previous
        completion. They're treated like the user's own cards -- model
        context, never trimmed -- but returned with their original source and
        locked=True. Rerolling = same request with the cards to keep locked.

        `exclude`: names the model must never suggest (the user's ban list).
        Matched by exact name (case and
        punctuation ignored, no fuzzy matching, so a typo can't ban the wrong
        card); unknown names are skipped with a warning. A banned card that's
        also in the user's own list or locked is kept, with a warning."""
        from training.ensemble import ensemble_complete_deck

        warnings: list[str] = []
        commander_names = [n.strip() for n in commander_names if n.strip()]
        if not 1 <= len(commander_names) <= 2:
            raise DeckInputError("Choose one commander, or two for a partner pair.")

        # -- commanders --
        commanders: list[CardRecord] = []
        unresolved: list[dict] = []
        for name in commander_names:
            card = self._resolve(name, warnings)
            if card is None:
                unresolved.append({"name": name, "line": None, "suggestions": self._suggestions(name, commander_only=True)})
            else:
                commanders.append(card)
        if unresolved:
            raise DeckInputError("Couldn't find that commander.", unresolved)
        too_new = [c.name for c in commanders if not self.model_knows(c)]
        if too_new:
            raise DeckInputError(f"{' and '.join(too_new)} is newer than the model, so it can't build around it yet.")
        if len({c.oracle_id for c in commanders}) != len(commanders):
            raise DeckInputError("The two commanders are the same card.")
        for card in commanders:
            if not is_commander_eligible(card):
                warnings.append(f"{card.name} can't normally be a commander.")
            if card.commander_legality == "banned":
                warnings.append(f"{card.name} is banned in Commander.")
        if len(commanders) == 2 and not _partner_pair_ok(*commanders):
            warnings.append(f"{commanders[0].name} and {commanders[1].name} don't look like a legal partner pair.")
        colors: set[str] = set()
        for card in commanders:
            colors |= set(card.color_identity)

        # -- decklist --
        known: dict[str, tuple[CardRecord, int]] = {}  # oracle_id -> (card, qty), input order
        basics: Counter = Counter()
        commander_ids = {c.oracle_id for c in commanders}
        for entry in parse_decklist(decklist):
            basic = next((b for b in BASIC_LAND_NAMES if b.lower() == entry.name.lower()), None)
            if basic:
                basics[basic] += entry.quantity
                continue
            card = self._resolve(entry.name, warnings)
            if card is None:
                unresolved.append({"name": entry.name, "line": entry.line_number, "suggestions": self._suggestions(entry.name)})
                continue
            if card.oracle_id in commander_ids:
                continue  # exports often list the commander inside the deck too
            _, qty = known.get(card.oracle_id, (card, 0))
            known[card.oracle_id] = (card, qty + entry.quantity)
        if unresolved:
            noun = "card" if len(unresolved) == 1 else "cards"
            raise DeckInputError(f"Couldn't find {len(unresolved)} {noun} in the card database.", unresolved)

        # -- locked cards from a previous completion --
        # Each entry is (name, quantity, source, is_locked); is_locked=False
        # marks a card the reroll keeps without the user locking it. Only
        # locked cards are safe from being cut, and at most 100 minus the
        # commanders can be locked, so a deck can always be cut to 100.
        locked_sources: dict[str, str] = {}  # oracle_id -> original source
        protected: set[str] = set()  # oracle_ids the user locked
        locked_basics: list[tuple[str, int, str, bool]] = []
        lock_cap = DECK_SIZE - len(commanders)
        locked_count = sum(entry[1] for entry in locked or [] if len(entry) < 4 or entry[3])
        if locked_count > lock_cap:
            raise DeckInputError(f"You can lock at most {lock_cap} cards (you locked {locked_count}).")
        for name, qty, source, *flag in locked or []:
            is_locked = flag[0] if flag else True
            basic = next((b for b in BASIC_LAND_NAMES if b.lower() == name.strip().lower()), None)
            if basic:
                locked_basics.append((basic, qty, source, is_locked))
                continue
            card = self._resolve(name, [])
            if card is None:
                raise DeckInputError(f"Couldn't find locked card {name!r}.")
            if card.oracle_id in commander_ids or card.oracle_id in known:
                continue  # already in the deck as a commander or your own card
            known[card.oracle_id] = (card, qty)
            locked_sources[card.oracle_id] = source
            if is_locked:
                protected.add(card.oracle_id)

        for oid, (card, qty) in list(known.items()):
            if qty > 1 and not _allows_multiple(card):
                warnings.append(f"Commander is singleton: kept 1 copy of {card.name} (you listed {qty}).")
                known[oid] = (card, 1)
            if not set(card.color_identity) <= colors:
                warnings.append(f"{card.name} is outside your commander's color identity.")
            if card.commander_legality == "banned":
                warnings.append(f"{card.name} is banned in Commander.")
        if target_lands is None:
            target_lands = DEFAULT_TARGET_LANDS
        if target_lands < MIN_TARGET_LANDS:
            raise DeckInputError(f"Decks need at least {MIN_TARGET_LANDS} lands.")

        # -- over 100: cut down to size (locked and user cards included:
        # that's all an over-100 deck can be made of) --
        size = len(commanders) + sum(q for _, q in known.values()) + sum(basics.values())
        size += sum(q for _, q, _, _ in locked_basics)
        trimmed = False
        if size > DECK_SIZE:
            with self._lock, torch.no_grad():
                cut, left_over = self._trim_to_size(commanders, colors, known, protected, basics, locked_basics,
                                                    size - DECK_SIZE, target_lands)
            trimmed = True
            listing = ", ".join(f"{q} {name}" if q > 1 else name for name, q in cut.items())
            warnings.append(f"Your deck had {size} cards, so {size - DECK_SIZE - left_over} were cut "
                            f"to bring it to {DECK_SIZE}: {listing}.")
            if left_over:  # only reachable through the API with nothing left to cut
                warnings.append(f"It's still {left_over} over {DECK_SIZE}: every card left is locked.")

        unknown_to_model = [card.name for card, _ in known.values() if not self.model_knows(card)]
        if unknown_to_model:
            warnings.append(
                f"The model hasn't seen {', '.join(unknown_to_model)} yet (newer than its training data). "
                f"{'It stays' if len(unknown_to_model) == 1 else 'They stay'} in your deck, "
                f"but the picks don't take {'it' if len(unknown_to_model) == 1 else 'them'} into account."
            )
        for basic in basics:
            if basic != "Wastes" and BASIC_COLOR[basic] not in colors:
                warnings.append(f"{basic} is outside your commander's color identity.")

        input_cards = []
        for oid, (card, qty) in known.items():
            deck_card = self._deck_card(card, qty, locked_sources.get(oid, "input"), None)
            deck_card.locked = oid in protected
            input_cards.append(deck_card)
        input_cards += [self._basic_card(name, qty, "input") for name, qty in basics.items()]
        for name, qty, source, is_locked in locked_basics:
            deck_card = self._basic_card(name, qty, source)
            deck_card.locked = is_locked
            input_cards.append(deck_card)
        commander_cards = [self._deck_card(card, 1, "input", None) for card in commanders]

        # -- ban list --
        exclude_ids: list[int] = []
        unknown_excluded: list[str] = []
        seen_tokens: set[str] = set()
        for name in exclude or []:
            token = self.tok._name_index.get(_normalize_name(name))
            if token is None:
                if name not in unknown_excluded:
                    unknown_excluded.append(name)
                continue
            if token in seen_tokens:
                continue
            seen_tokens.add(token)
            card = self.tok.card_by_token[token]
            if card.oracle_id in known or card.oracle_id in commander_ids:
                warnings.append(f"{card.name} is on your ban list, but it's in your deck, so it was kept.")
                continue
            exclude_ids.append(self.tok.token_to_id[token])
        if unknown_excluded:
            warnings.append(f"Ignored unknown cards on your ban list: {', '.join(unknown_excluded)}.")

        self.prices.maybe_refresh()
        result = Completion(commanders=commander_cards, cards=input_cards, warnings=warnings,
                            target_lands=target_lands, prices_updated=self.prices.built_at)

        open_slots = DECK_SIZE - len(commanders) - sum(c.quantity for c in input_cards)
        if open_slots <= 0:  # (below 0 only when trimming had to stop at the land target)
            if not trimmed:
                warnings.append(f"Your list is already {DECK_SIZE} cards, so nothing was added.")
            return self._warn_if_short(result)

        # -- model picks, one per open slot. A fresh seed every request, so
        # submitting the same inputs again gives a different deck. --
        seed = secrets.randbelow(2**31)
        with self._lock, torch.no_grad():
            picks = ensemble_complete_deck(
                [c.name for c in commanders], [c.name for c, _ in known.values() if self.model_knows(c)], open_slots,
                self.tok, self.device, self.staple_model, self.synergy_model, self.tiers,
                prune_model=self.prune_model if prune else None,
                chunk_size=GENERATION_CHUNK_SIZE,
                seed=seed, allow_banned=False, exclude_ids=exclude_ids,
            )
        pick_cards = [(self._card(name), score, source) for name, score, source in picks]

        # -- land target: swap the weakest nonland picks for basics --
        land_count = sum(qty for card, qty in known.values() if is_front_land(card)) + sum(basics.values())
        land_count += sum(qty for _, qty, _, _ in locked_basics)
        land_count += sum(is_front_land(card) for card, _, _ in pick_cards)
        basics_needed = max(0, target_lands - land_count)
        if basics_needed:
            removable = [
                i for i, (card, _, source) in enumerate(pick_cards)
                if not is_front_land(card) and source == "synergy"
                and self.tiers.get(self.tok.token_to_id[self.tok.oracle_id_to_token[card.oracle_id]]) != "staple"
            ]
            removable.sort(key=lambda i: pick_cards[i][1])
            drop = set(removable[:basics_needed])
            basics_needed = len(drop)
            pick_cards = [p for i, p in enumerate(pick_cards) if i not in drop]

        result.cards += [self._deck_card(card, 1, source, score) for card, score, source in pick_cards]

        if basics_needed:
            pips: Counter = Counter()
            for card in commanders + [c for c, _ in known.values()] + [c for c, _, _ in pick_cards]:
                if not is_front_land(card):
                    for color, n in zip(COLORS, _pip_counts(card.mana_cost)):
                        pips[color] += float(n)
            for name, qty in allocate_basics(basics_needed, colors, pips).items():
                result.cards.append(self._basic_card(name, qty, "basic"))

        return self._warn_if_short(result)

    def _trim_to_size(
        self,
        commanders: list[CardRecord],
        colors: set[str],
        known: dict[str, tuple[CardRecord, int]],
        protected: set[str],
        basics: Counter,
        locked_basics: list[tuple[str, int, str, bool]],
        excess: int,
        target_lands: int,
    ) -> tuple[Counter, int]:
        """Cuts `excess` cards from an over-100 deck, editing `known`,
        `basics` and `locked_basics` in place. Returns (what was cut, by
        name; how many it couldn't cut). Locked cards (`protected`, and
        locked basics) are never cut. Among the rest, in order:

        1. Cards that make the deck illegal: outside the commander's color
           identity, or banned in Commander.
        2. Basic lands above the land target, from the biggest pile.
        3. Nonland, non-staple cards, by the prune model's fit score (lowest
           first, re-scored after every cut, as the ensemble's own pruning
           does). Staples are held back until after these: the prune model
           is known to doubt near-universal cards wrongly.
        4. Staples, by the same score.
        5. Nonbasic lands above the target, by the same score.
        6. Cards the model doesn't know yet (no score), last listed first.
        7. Only if still over (the locked cards leave nothing else): lands
           below the target, basics first. Locks win over the land target.
        """
        cut: Counter = Counter()

        def land_count() -> int:
            return (sum(q for c, q in known.values() if is_front_land(c)) + sum(basics.values())
                    + sum(q for _, q, _, _ in locked_basics))

        def token_id(card: CardRecord) -> int:
            return self.tok.token_to_id[self.tok.oracle_id_to_token[card.oracle_id]]

        def cuttable() -> list[str]:
            return [oid for oid in known if oid not in protected]

        def drop_one(oid: str) -> None:
            card, qty = known[oid]
            if qty > 1:
                known[oid] = (card, qty - 1)
            else:
                del known[oid]
            cut[card.name] += 1

        def lands_spare(card: CardRecord) -> bool:
            return not is_front_land(card) or land_count() > target_lands

        def cut_basic(above_target: bool) -> bool:
            """Cuts one unlocked basic from the biggest pile; False if none can go."""
            if above_target and land_count() <= target_lands:
                return False
            piles = Counter(basics)
            for name, qty, _, is_locked in locked_basics:
                if not is_locked:
                    piles[name] += qty
            if not piles:
                return False
            name = max(piles, key=piles.get)
            if basics[name]:
                basics[name] -= 1
                if not basics[name]:
                    del basics[name]
            else:
                i = next(i for i, (n, _, _, lk) in enumerate(locked_basics) if n == name and not lk)
                n, qty, source, is_locked = locked_basics[i]
                if qty > 1:
                    locked_basics[i] = (n, qty - 1, source, is_locked)
                else:
                    del locked_basics[i]
            cut[name] += 1
            return True

        # 1. illegal cards
        for oid in cuttable():
            card = known[oid][0]
            illegal = not set(card.color_identity) <= colors or card.commander_legality == "banned"
            while excess and illegal and oid in known and lands_spare(card):
                drop_one(oid)
                excess -= 1

        # 2. basics above the land target, biggest pile first
        while excess and cut_basic(above_target=True):
            excess -= 1

        # 3-5. by the prune model's score, re-scored after every cut
        cmdr_ids = [token_id(c) for c in commanders]
        cmdr_tensor = torch.tensor([cmdr_ids], dtype=torch.long, device=self.device)
        cmdr_mask = torch.ones_like(cmdr_tensor, dtype=torch.bool)

        def is_staple(card: CardRecord) -> bool:
            return self.tiers.get(token_id(card)) == "staple"

        stages = [
            lambda c: not is_front_land(c) and not is_staple(c),
            lambda c: not is_front_land(c) and is_staple(c),
            lambda c: is_front_land(c) and land_count() > target_lands,
            lambda c: is_front_land(c),  # 7, if the locks leave nothing else
        ]
        for n, eligible in enumerate(stages):
            if n == 3:
                # 6. cards the model can't score, last listed first
                for oid in reversed(cuttable()):
                    card = known[oid][0]
                    while excess and oid in known and not self.model_knows(card) and lands_spare(card):
                        drop_one(oid)
                        excess -= 1
                # 7. basics, then other lands, even below the target
                while excess and cut_basic(above_target=False):
                    excess -= 1
            while excess:
                scored = [oid for oid, (c, _) in known.items() if self.model_knows(c)]
                candidates = [oid for oid in scored if oid not in protected and eligible(known[oid][0])]
                if not candidates:
                    break
                deck_ids = [token_id(known[oid][0]) for oid in scored]
                deck_tensor = torch.tensor([deck_ids], dtype=torch.long, device=self.device)
                logits = self.prune_model(deck_tensor, torch.ones_like(deck_tensor, dtype=torch.bool),
                                          cmdr_tensor, cmdr_mask)[0].cpu()
                fit = dict(zip(scored, torch.sigmoid(logits).tolist()))
                drop_one(min(candidates, key=fit.__getitem__))
                excess -= 1
        return cut, excess

    @staticmethod
    def _warn_if_short(result: Completion) -> Completion:
        """Basics only replace the model's own new picks, never the user's or
        locked cards, so a deck made mostly of those can end up short."""
        lands = sum(c.quantity for c in result.cards if c.category == "Land")
        if lands < result.target_lands:
            result.warnings.append(
                f"This deck has {lands} lands, below the {result.target_lands}-land target, because your own and "
                "locked cards fill the slots basics would need. Add lands to your list or unlock some cards."
            )
        return result