recomb.py

"""Recombinator odds calculator.

Finds optimal crafting recipes for recombining items, ranked by
lowest expected cost (fewest bases consumed on average).
"""

from __future__ import annotations

import argparse
import re
from collections import namedtuple
from dataclasses import dataclass, field
from itertools import product


# ---------------------------------------------------------------------------
# Probability table
# PROB[pool_size][final_count] = probability
# Pool size = sum of mods from both input items (for one mod type)
# Final count = number of mods on the output item
# ---------------------------------------------------------------------------
PROB: dict[int, dict[int, float]] = {
    1: {0: 0.41, 1: 0.59, 2: 0.00, 3: 0.00},
    2: {0: 0.00, 1: 0.67, 2: 0.33, 3: 0.00},
    3: {0: 0.00, 1: 0.39, 2: 0.52, 3: 0.10},
    4: {0: 0.00, 1: 0.11, 2: 0.59, 3: 0.31},
    5: {0: 0.00, 1: 0.00, 2: 0.43, 3: 0.57},
    6: {0: 0.00, 1: 0.00, 2: 0.28, 3: 0.72},
}


def pool_prob(pool_size: int, final_count: int) -> float:
    """Probability of getting `final_count` mods from a pool of `pool_size`."""
    if pool_size == 0:
        return 1.0 if final_count == 0 else 0.0
    if pool_size < 0 or pool_size > 6:
        return 0.0
    if final_count < 0 or final_count > 3:
        return 0.0
    return PROB[pool_size][final_count]


# ---------------------------------------------------------------------------
# Data model
# ---------------------------------------------------------------------------
Item = namedtuple("Item", ["p", "s"])


def item_str(item: Item) -> str:
    return f"{item.p}p{item.s}s"


def parse_item(text: str) -> Item:
    """Parse '3p1s' or '3p 1s' style strings into an Item."""
    m = re.match(r"(\d+)\s*p\s*(\d+)\s*s", text.strip().lower())
    if not m:
        raise ValueError(f"Cannot parse item: {text!r}  (expected e.g. '3p1s')")
    return Item(int(m.group(1)), int(m.group(2)))


BASE_ITEMS = {Item(1, 0), Item(0, 1)}


@dataclass
class Recipe:
    target: Item
    prob: float  # single-attempt probability of entire tree succeeding
    expected_cost: float  # expected cost including retries
    step_prob: float  # probability of just this recombination step


@dataclass
class BaseRecipe(Recipe):
    """Acquiring a base item — no crafting needed."""
    pass


@dataclass
class CombineRecipe(Recipe):
    """A recombination step combining two sub-recipes."""
    left: Recipe = field(default=None)
    right: Recipe = field(default=None)


@dataclass
class InventoryRecipe(Recipe):
    """An item the player already has — free, infinite supply."""
    pass


# ---------------------------------------------------------------------------
# Core helpers
# ---------------------------------------------------------------------------
def recomb_step_prob(left: Item, right: Item, target: Item) -> float:
    """Probability of a single recombination step producing `target`."""
    pp = pool_prob(left.p + right.p, target.p)
    ps = pool_prob(left.s + right.s, target.s)
    return pp * ps


def expected_bases(recipe: Recipe) -> float:
    """Expected number of base items consumed to produce one copy of this recipe."""
    if isinstance(recipe, InventoryRecipe):
        return 0.0
    if isinstance(recipe, BaseRecipe):
        return 1.0
    if isinstance(recipe, CombineRecipe):
        return (expected_bases(recipe.left) + expected_bases(recipe.right)) / recipe.step_prob
    return 0.0


def expected_recombs(recipe: Recipe, inventory: set[Item] | None = None) -> float:
    """Expected number of recombinator uses to produce one copy of this recipe.

    If inventory is provided, up to 1 copy of each inventory item can be skipped.
    """
    if isinstance(recipe, (BaseRecipe, InventoryRecipe)):
        return 0.0

    if not inventory:
        # No inventory, count everything
        if isinstance(recipe, CombineRecipe):
            return (expected_recombs(recipe.left) + expected_recombs(recipe.right) + 1) / recipe.step_prob
        return 0.0

    # Track how many times we've used each inventory item (limit: 1 per item)
    used: dict[Item, int] = {}

    def calc(r: Recipe) -> float:
        if isinstance(r, (BaseRecipe, InventoryRecipe)):
            return 0.0
        if isinstance(r, CombineRecipe):
            # Check if we can skip this step using inventory
            can_skip = (r.target in inventory and used.get(r.target, 0) < 1)
            if can_skip:
                used[r.target] = used.get(r.target, 0) + 1
                return 0.0
            return (calc(r.left) + calc(r.right) + 1) / r.step_prob
        return 0.0

    return calc(recipe)


def count_inventory_available(recipe: Recipe, inventory: set[Item] | None = None) -> dict[Item, int]:
    """Count how many times each inventory item is needed as an intermediate target.

    Returns the number of times we could skip a step for each inventory item.
    Caller should use min(count, 1) to respect the 1-copy limit.
    """
    if not inventory:
        return {}

    needed: dict[Item, int] = {}

    def walk(r: Recipe) -> None:
        if isinstance(r, CombineRecipe):
            # If this step's target is in inventory, we could potentially skip it
            if r.target in inventory:
                needed[r.target] = needed.get(r.target, 0) + 1
            # Continue walking the tree
            walk(r.left)
            walk(r.right)

    walk(recipe)
    return needed


# ---------------------------------------------------------------------------
# DP Optimizer — top-K recipes
# ---------------------------------------------------------------------------
def find_best_recipes(
    target: Item,
    top_k: int = 5,
    base_cost: float = 1.0,
    dust_cost: float = 1.0,
) -> list[Recipe]:
    """Find the top-K cheapest recipes for `target`.

    Returns a list of up to `top_k` Recipe objects, sorted by expected cost
    (lowest first).

    Uses bottom-up DP: items are processed in order of ascending total mods
    (p+s), so all possible inputs are already solved before each item.
    """
    memo: dict[Item, list[Recipe]] = {}

    # Seed base items
    for bi in BASE_ITEMS:
        base = BaseRecipe(target=bi, prob=1.0, expected_cost=base_cost, step_prob=1.0)
        memo[bi] = [base]

    # Process all items bottom-up by total mods
    all_items = [Item(p, s) for p in range(4) for s in range(4) if p + s > 0]
    all_items.sort(key=lambda i: (i.p + i.s, i.p, i.s))

    for item in all_items:
        if item in memo:
            continue  # Already set as base item

        candidates: list[Recipe] = []

        for p1, s1, p2, s2 in product(range(4), range(4), range(4), range(4)):
            left_item = Item(p1, s1)
            right_item = Item(p2, s2)
            if left_item == item or right_item == item:
                continue
            if (p1 + s1) == 0 or (p2 + s2) == 0:
                continue
            # Canonical ordering to avoid duplicates
            if left_item > right_item:
                continue

            sp = recomb_step_prob(left_item, right_item, item)
            if sp <= 0:
                continue

            left_recipes = memo.get(left_item, [])
            right_recipes = memo.get(right_item, [])

            for lr in left_recipes:
                for rr in right_recipes:
                    ec = (lr.expected_cost + rr.expected_cost + dust_cost) / sp
                    total_prob = sp * lr.prob * rr.prob
                    candidates.append(
                        CombineRecipe(
                            target=item,
                            prob=total_prob,
                            expected_cost=ec,
                            step_prob=sp,
                            left=lr,
                            right=rr,
                        )
                    )

        memo[item] = _top_k_unique(
            candidates, top_k, key=lambda r: r.expected_cost, reverse=False
        )

    return memo.get(target, [])


def _recipe_signature(recipe: Recipe) -> str:
    """Generate a structural signature for deduplication."""
    if isinstance(recipe, InventoryRecipe):
        return f"I({item_str(recipe.target)})"
    if isinstance(recipe, BaseRecipe):
        return f"B({item_str(recipe.target)})"
    if isinstance(recipe, CombineRecipe):
        left_sig = _recipe_signature(recipe.left)
        right_sig = _recipe_signature(recipe.right)
        # Canonical order
        if left_sig > right_sig:
            left_sig, right_sig = right_sig, left_sig
        return f"C({left_sig}+{right_sig}->{item_str(recipe.target)})"
    return str(id(recipe))


def _top_k_unique(candidates: list[Recipe], k: int, key, reverse: bool) -> list[Recipe]:
    """Return top-K unique recipes sorted by `key`."""
    seen: set[str] = set()
    result: list[Recipe] = []
    for r in sorted(candidates, key=key, reverse=reverse):
        sig = _recipe_signature(r)
        if sig not in seen:
            seen.add(sig)
            result.append(r)
            if len(result) >= k:
                break
    return result


# ---------------------------------------------------------------------------
# Output formatting
# ---------------------------------------------------------------------------
def format_recipe(recipe: Recipe) -> str:
    """Format a recipe tree as a human-readable string."""
    lines: list[str] = []
    steps: list[str] = []
    _collect_steps(recipe, steps)

    bases = expected_bases(recipe)
    lines.append(f"Target: {item_str(recipe.target)}")
    lines.append(f"Expected bases: ~{bases:.1f}")
    lines.append(f"Expected cost: {recipe.expected_cost:.2f} units")
    lines.append("")

    for i, step in enumerate(steps, 1):
        lines.append(f"  Step {i}: {step}")

    return "\n".join(lines)


def _collect_steps(recipe: Recipe, steps: list[str]) -> None:
    """Collect recombination steps in bottom-up order."""
    if isinstance(recipe, (BaseRecipe, InventoryRecipe)):
        return
    if isinstance(recipe, CombineRecipe):
        _collect_steps(recipe.left, steps)
        _collect_steps(recipe.right, steps)
        left_str = item_str(recipe.left.target)
        right_str = item_str(recipe.right.target)
        target_str = item_str(recipe.target)
        pp = recipe.left.target.p + recipe.right.target.p
        sp = recipe.left.target.s + recipe.right.target.s
        p_prefix = pool_prob(pp, recipe.target.p)
        p_suffix = pool_prob(sp, recipe.target.s)
        steps.append(
            f"{left_str} + {right_str} -> {target_str}  "
            f"(prefix {pp}->{recipe.target.p}: {p_prefix*100:.1f}%, "
            f"suffix {sp}->{recipe.target.s}: {p_suffix*100:.1f}%, "
            f"combined: {recipe.step_prob*100:.2f}%)"
        )


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------
def main():
    parser = argparse.ArgumentParser(description="Recombinator odds calculator")
    parser.add_argument("target", help="Target item, e.g. '3p1s'")
    parser.add_argument("--top-k", type=int, default=5, help="Number of recipes to show (default: 5)")
    parser.add_argument("--base-cost", type=float, default=1.0, help="Cost of a base item (default: 1.0)")
    parser.add_argument("--dust-cost", type=float, default=1.0, help="Cost of dust per recombination (default: 1.0)")
    parser.add_argument("--inventory", "-i", nargs="*", default=[], help="Items you already have, e.g. '2p1s 1p2s'")
    args = parser.parse_args()

    target = parse_item(args.target)
    if target.p < 0 or target.p > 3 or target.s < 0 or target.s > 3:
        parser.error("Prefixes and suffixes must be 0-3")

    inv = {parse_item(x) for x in args.inventory} if args.inventory else None
    recipes = find_best_recipes(
        target, top_k=args.top_k, base_cost=args.base_cost, dust_cost=args.dust_cost,
    )

    print("=" * 60)
    print(f"Recipinator — Target: {item_str(target)}")
    print(f"Base cost: {args.base_cost}, Dust cost: {args.dust_cost}")
    if inv:
        print(f"Inventory: {', '.join(item_str(i) for i in sorted(inv))}")
    print("=" * 60)

    for i, recipe in enumerate(recipes, 1):
        print(f"\nRecipe #{i}:")
        recombs = expected_recombs(recipe, inv)
        print(f"Expected recombinator uses: ~{recombs:.0f}")
        # Show how many times each inventory item is needed
        if inv:
            needed = count_inventory_available(recipe, inv)
            if needed:
                for item, count in sorted(needed.items()):
                    available = 1  # Each toggle = 1 copy
                    used = min(count, available)
                    status = f"using {used}/{count} occurrence{'s' if count > 1 else ''}"
                    print(f"  {item_str(item)}: {status}")
        print(format_recipe(recipe))


if __name__ == "__main__":
    main()