Balatro probabilities

import random
from typing import List, Tuple, NamedTuple


class Card(NamedTuple):
    suit: str
    rank: int


def create_deck() -> List[Card]:
    """
    Creates a standard 52-card deck.
    """
    suits = ["S", "H", "D", "C"]
    ranks = list(range(2, 15))
    return [Card(suit, rank) for suit in suits for rank in ranks]


def draw_cards(deck: List[Card], num_cards: int) -> Tuple[List[Card], List[Card]]:
    """
    Draws a specified number of cards from the deck and returns the drawn cards
    and the updated deck.
    """
    return random.sample(deck, num_cards)


def contains_two_pair(hand: List[Card]) -> bool:
    rank_counts = {}
    for card in hand:
        rank_counts[card.rank] = rank_counts.get(card.rank, 0) + 1
    pairs = 0
    for count in rank_counts.values():
        if count == 2:
            pairs += 1
    return pairs >= 2


def contains_three_of_a_kind(hand: List[Card]) -> bool:
    rank_counts = {}
    for card in hand:
        rank_counts[card.rank] = rank_counts.get(card.rank, 0) + 1
    values = rank_counts.values()
    sorted_values = sorted(values, reverse=True)
    # if the second value is >1 then we have a full-house
    return sorted_values[0] == 3 and sorted_values[1] == 1


def contains_flush(hand: List[Card]) -> bool:
    suit_counts = {}
    for card in hand:
        suit_counts[card.suit] = suit_counts.get(card.suit, 0) + 1
    return max(suit_counts.values()) >= 5


def contains_straight(hand: List[Card]) -> bool:
    ranks = {card.rank for card in hand}
    if 14 in ranks:  # Ace can be high or low, so add 1 to ranks if Ace is present
        ranks.add(1)
    for start in range(1, 10):  # 10, J, Q, K, A (high) is the highest starting point
        if all(rank in ranks for rank in range(start, start + 5)):
            return True
    return False


def contains_full_house(hand: List[Card]) -> bool:
    rank_counts = {}
    for card in hand:
        rank_counts[card.rank] = rank_counts.get(card.rank, 0) + 1
    values = rank_counts.values()
    sorted_values = sorted(values, reverse=True)[:2]
    return sorted_values[0] == 3 and sorted_values[1] >= 2
    return


def contains_four_of_a_kind(hand: List[Card]) -> bool:
    rank_counts = {}
    for card in hand:
        rank_counts[card.rank] = rank_counts.get(card.rank, 0) + 1
    return 4 in rank_counts.values()


def simulate_scenario(scenario_func: callable, num_simulations: int) -> float:
    """
    Simulates a scenario function for a given number of simulations and returns
    the probability of success.
    """
    successes = 0
    for _ in range(num_simulations):
        if scenario_func():
            successes += 1
    return successes / num_simulations


def scenario_1() -> bool:
    """
    Scenario 1: Starting with 4 cards of the same suit from an 8-card hand,
    discarding the other 4 cards, redrawing 4 cards, and checking if the
    resulting hand has a flush.
    """
    deck = create_deck()
    hand = [
        Card("S", 14),
        Card("S", 13),
        Card("S", 12),
        Card("S", 11),
        Card("D", 2),
        Card("D", 3),
        Card("H", 2),
        Card("C", 2),
    ]
    remaining_deck = [card for card in deck if card not in hand]
    new_hand = hand[:4] + draw_cards(remaining_deck, 4)
    return contains_flush(new_hand)


def scenario_2() -> bool:
    """
    Scenario 2: Starting with two pairs in an 8 card hand, discarding the other
    4 cards, and checking if the resulting hand has a full house.
    """
    deck = create_deck()
    hand = [
        Card("S", 14),
        Card("H", 14),
        Card("S", 7),
        Card("H", 7),
        Card("D", 2),
        Card("D", 3),
        Card("C", 4),
        Card("C", 5),
    ]
    remaining_deck = [card for card in deck if card not in hand]
    new_hand = hand[:4] + draw_cards(remaining_deck, 4)
    return contains_full_house(new_hand)


def scenario_3() -> bool:
    """
    Scenario 3: Starting with two pairs in an 8 card hand, discarding the other
    4 cards, and checking if the resulting hand has a four-of-a-kind.
    """
    deck = create_deck()
    hand = [
        Card("S", 14),
        Card("H", 14),
        Card("S", 7),
        Card("H", 7),
        Card("D", 2),
        Card("D", 3),
        Card("C", 4),
        Card("C", 5),
    ]
    remaining_deck = [card for card in deck if card not in hand]
    new_hand = hand[:4] + draw_cards(remaining_deck, 4)
    return contains_four_of_a_kind(new_hand)


def scenario_4() -> bool:
    """
    Scenario 4: Drawing an open-ended straight (e.g., 5-6-7-8) from an 8-card hand, discarding the other 4 cards,
    redrawing 4 cards, and checking if the resulting hand has a straight.
    """
    deck = create_deck()
    hand = [
        Card("S", 5),
        Card("S", 6),
        Card("H", 7),
        Card("C", 8),
        Card("D", 2),
        Card("D", 3),
        Card("C", 10),
        Card("H", 10),
    ]
    remaining_deck = [card for card in deck if card not in hand]
    new_hand = hand[:4] + draw_cards(remaining_deck, 4)
    return contains_straight(new_hand)


def scenario_5() -> bool:
    """
    Scenario 5: Drawing a pair from an 8-card hand, discarding 5 cards,
    redrawing 5 cards, and checking if the resulting hand has a three of a kind.
    """
    deck = create_deck()
    hand = [
        Card("S", 4),
        Card("H", 4),
        Card("D", 5),
        Card("C", 6),
        Card("S", 7),
        Card("H", 8),
        Card("D", 9),
        Card("C", 10),
    ]
    remaining_deck = [card for card in deck if card not in hand]
    new_hand = hand[:2] + draw_cards(remaining_deck, 6)
    return contains_three_of_a_kind(new_hand)


def scenario_6() -> bool:
    """
    Scenario 6: Drawing a pair from an 8-card hand, discarding the other 5 cards,
    redrawing 5 cards, and checking if the resulting hand has two pairs.
    """
    deck = create_deck()
    hand = [
        Card("S", 4),
        Card("H", 4),
        Card("D", 5),
        Card("C", 6),
        Card("S", 7),
        Card("H", 8),
        Card("D", 9),
        Card("C", 10),
    ]
    remaining_deck = [card for card in deck if card not in hand]
    new_hand = hand[:3] + draw_cards(remaining_deck, 5)
    return contains_two_pair(new_hand)


def scenario_7() -> bool:
    """
    Scenario 7: Drawing a pair from an 8-card hand, discarding five cards,
    redrawing five cards, and checking if the resulting hand has a full house.
    """
    deck = create_deck()
    hand = [
        Card("S", 4),
        Card("H", 4),
        Card("D", 5),
        Card("C", 6),
        Card("S", 7),
        Card("H", 8),
        Card("D", 9),
        Card("C", 10),
    ]
    remaining_deck = [card for card in deck if card not in hand]
    new_hand = hand[:3] + draw_cards(remaining_deck, 5)
    return contains_full_house(new_hand)


def scenario_8() -> bool:
    """
    Scenario 8: Drawing a pair from an 8-card hand, discarding five cards,
    redrawing five cards, and checking if the resulting hand has a four-of-a-kind.
    """
    deck = create_deck()
    hand = [
        Card("S", 4),
        Card("H", 4),
        Card("D", 5),
        Card("C", 6),
        Card("S", 7),
        Card("H", 8),
        Card("D", 9),
        Card("C", 10),
    ]
    remaining_deck = [card for card in deck if card not in hand]
    new_hand = hand[:3] + draw_cards(remaining_deck, 5)
    return contains_four_of_a_kind(new_hand)


num_simulations = 1000000

print("Scenario 1 probability:", simulate_scenario(scenario_1, num_simulations))
print("Scenario 2 probability:", simulate_scenario(scenario_2, num_simulations))
print("Scenario 3 probability:", simulate_scenario(scenario_3, num_simulations))
print("Scenario 4 probability:", simulate_scenario(scenario_4, num_simulations))
print("Scenario 5 probability:", simulate_scenario(scenario_5, num_simulations))
print("Scenario 6 probability:", simulate_scenario(scenario_6, num_simulations))
print("Scenario 7 probability:", simulate_scenario(scenario_7, num_simulations))
print("Scenario 8 probability:", simulate_scenario(scenario_8, num_simulations))

➜ python3 balatro.py
Scenario 1 probability: 0.61475
Scenario 2 probability: 0.322597
Scenario 3 probability: 0.012559
Scenario 4 probability: 0.565939
Scenario 5 probability: 0.11537
Scenario 6 probability: 0.525783
Scenario 7 probability: 0.155742
Scenario 8 probability: 0.011344