import random
import pandas as pd
from collections import Counter

def can_form_set(cnt):
    jokers = cnt.get('J', 0)
    # Check three of a kind with jokers
    for t in ['I', 'C', 'A']:
        if cnt.get(t, 0) + jokers >= 3:
            return True
    # Check mixed set with jokers: need 3 distinct types
    distinct = sum(1 for t in ['I', 'C', 'A'] if cnt.get(t, 0) >= 1)
    if distinct + jokers >= 3:
        return True
    return False

def simulate_once(heuristic):
    # Initialize and shuffle deck once
    deck = ['I'] * 14 + ['C'] * 14 + ['A'] * 14 + ['J'] * 2
    random.shuffle(deck)
    hand = []

    events_3 = success_3 = events_4 = success_4 = 0
    kind_trades = mixed_trades = 0

    while deck:
        while len(hand) < 5 and deck:
            hand.append(deck.pop(random.randrange(len(deck))))
            cnt = Counter(hand)
            if len(hand) == 3:
                events_3 += 1
                if can_form_set(cnt): success_3 += 1
            if len(hand) == 4:
                events_4 += 1
                if can_form_set(cnt): success_4 += 1

        if len(hand) < 5:
            break

        cnt = Counter(hand)
        possible_sets = []
        for t in ['I', 'C', 'A']:
            if cnt[t] + cnt['J'] >= 3:
                jokers_needed = max(0, 3 - cnt[t])
                possible_sets.append(('kind', t, jokers_needed))
        distinct = sum(1 for t in ['I', 'C', 'A'] if cnt[t] >= 1)
        if distinct + cnt['J'] >= 3:
            jokers_needed = max(0, 3 - distinct)
            possible_sets.append(('mixed', None, jokers_needed))

        # Determine sub-heuristics
        base, joker_pref = heuristic.split('_')
        # Filter by joker preference
        if joker_pref == 'first':
            with_joker = [s for s in possible_sets if s[2] > 0]
            p_sets = with_joker if with_joker else possible_sets
        else:  # avoid
            without_joker = [s for s in possible_sets if s[2] == 0]
            p_sets = without_joker if without_joker else possible_sets

        # Choose based on base heuristic
        if base == 'good':
            chosen = next((s for s in p_sets if s[0] == 'kind'), p_sets[0])
        else:  # bad
            chosen = next((s for s in p_sets if s[0] == 'mixed'), p_sets[0])

        trade_type, trade_val, jokers_needed = chosen
        if trade_type == 'kind':
            for _ in range(3 - jokers_needed):
                hand.remove(trade_val)
            for _ in range(jokers_needed):
                hand.remove('J')
            kind_trades += 1
        else:
            for t in ['I', 'C', 'A']:
                if t in hand:
                    hand.remove(t)
                else:
                    hand.remove('J')
            mixed_trades += 1

    return events_3, success_3, events_4, success_4, kind_trades, mixed_trades

def run_simulation(num_games):
    heuristics = ['good_first', 'good_avoid', 'bad_first', 'bad_avoid']
    records = []
    for heuristic in heuristics:
        total_e3 = total_s3 = total_e4 = total_s4 = 0
        total_kind = total_mixed = 0
        for _ in range(num_games):
            e3, s3, e4, s4, kt, mt = simulate_once(heuristic)
            total_e3 += e3; total_s3 += s3
            total_e4 += e4; total_s4 += s4
            total_kind += kt; total_mixed += mt
        records.append({
            'Heuristic': heuristic,
            'P(Tradeable at 3 cards)': total_s3 / total_e3 if total_e3 else 0,
            'P(Tradeable at 4 cards)': total_s4 / total_e4 if total_e4 else 0,
            'Avg Kind Trades per Game': total_kind / num_games,
            'Avg Mixed Trades per Game': total_mixed / num_games
        })
    return pd.DataFrame(records)

# Run and display
df = run_simulation(20000)
print(df)