strongest_dice.py

import random
import math
import itertools
import matplotlib.pyplot as plt
from fractions import Fraction
from collections import defaultdict


class Dice:
    def __init__(self, side_values: list[int]):
        self.side_values = side_values

    def __str__(self) -> str:
        values = " ".join(str(v) for v in self.side_values)
        return f"<{values}>"

    def roll(self):
        return random.choice(self.side_values)


class DiceEngine:
    def make_dice_side_variants(self, *, n_sides: int, budget: int) -> list[list[int]]:
        if n_sides == 1:
            return [[budget]]

        variants = []
        biggest_first_slot_budget = int(math.ceil(budget / n_sides))
        for first_slot_budget in range(biggest_first_slot_budget, budget + 1):
            for remainder in self.make_dice_side_variants(
                n_sides=n_sides - 1,
                budget=budget - first_slot_budget,
            ):
                if any(side > first_slot_budget for side in remainder):
                    continue
                variants.append([first_slot_budget] + remainder)

        return variants

    @staticmethod
    def get_strongest_dice(dice1: Dice, dice2: Dice):
        expected_winrates = []
        for side1 in dice1.side_values:
            expected_wins = Fraction(0)
            for side2 in dice2.side_values:
                if side1 > side2:
                    expected_wins += Fraction(1)
                elif side1 == side2:
                    expected_wins += Fraction("1/2")
            expected_winrates.append(Fraction(expected_wins, len(dice2.side_values)))
        total_expected_winrate = Fraction(
            sum(expected_winrates),
            len(expected_winrates),
        )
        if total_expected_winrate > Fraction("1/2"):
            return dice1
        elif total_expected_winrate < Fraction("1/2"):
            return dice2
        else:
            return None

    def get_dice_stats(self, dice_list):
        stats = defaultdict(lambda: {"wins": set(), "losses": set(), "ties": set()})
        for dice1, dice2 in itertools.combinations(dice_list, r=2):
            winner = engine.get_strongest_dice(dice1, dice2)
            if winner == dice1:
                stats[dice1]["wins"].add(dice2)
                stats[dice2]["losses"].add(dice1)
            elif winner == dice2:
                stats[dice1]["losses"].add(dice2)
                stats[dice2]["wins"].add(dice1)
            else:
                stats[dice1]["ties"].add(dice2)
                stats[dice2]["ties"].add(dice1)
        return sorted(
            stats.items(),
            key=lambda item: len(item[1]["wins"]) + len(item[1]["ties"]) * 0.5,
            reverse=True,
        )

    def print_dice_stats(self, stats):
        for dice, stats in stats:
            wins = len(stats["wins"])
            losses = len(stats["losses"])
            ties = len(stats["ties"])
            print(f"{dice}: {wins} wins, {losses} losses, {ties} ties")

    def show_dice_table(self, dice_list):
        cellText = []
        cellColours = []
        for dice1 in dice_list:
            cellTextRow = []
            cellColoursRow = []
            for dice2 in dice_list:
                winner = engine.get_strongest_dice(dice1, dice2)
                text = {
                    dice1: "Win",
                    dice2: "Loss",
                    None: "Tie",
                }[winner]
                colour = {
                    dice1: "g",
                    dice2: "r",
                    None: "y",
                }[winner]
                cellTextRow.append(text)
                cellColoursRow.append(colour)
            cellText.append(cellTextRow)
            cellColours.append(cellColoursRow)

        fig, ax = plt.subplots()

        # hide axes
        fig.patch.set_visible(False)
        ax.axis("off")
        ax.axis("tight")

        ax.table(
            cellText=cellText,
            cellColours=cellColours,
            rowLabels=dice_list,
            colLabels=dice_list,
            loc="center",
        )

        fig.tight_layout()

        plt.show()

    def show_dice_image(self, dice_list):
        image = []
        for dice1 in dice_list:
            row = []
            for dice2 in dice_list:
                winner = engine.get_strongest_dice(dice1, dice2)
                pixel = {
                    dice1: [0, 255, 0],
                    dice2: [255, 0, 0],
                    None: [255, 255, 0],
                }[winner]
                row.append(pixel)
            image.append(row)

        plt.imshow(image)
        plt.show()


if __name__ == "__main__":
    engine = DiceEngine()

    all_possible_dice = [
        Dice(side_values)
        for side_values in engine.make_dice_side_variants(n_sides=6, budget=21)
    ]

    top_10 = [dice for dice, stats in engine.get_dice_stats(all_possible_dice)[:10]]
    engine.show_dice_table(top_10)