import random
from collections import defaultdict

DECK = [number+color+shade+symbol 
        for number in '123' for color in 'RGP' for shade in '@O=' for symbol in 'OSD']

def is_set(cards):
    "Are these 3 cards a set? Yes if each attribute has 1 or 3 values."
    assert len(cards) == 3
    return all(len(set(card[attr] for card in cards)) in (1, 3)
               for attr in range(4))

def find_set(cards):
    "Return a set found from this collection of cards, if there is one."
    for c in combinations(cards):
        if is_set(c):
            return c
    return []

def combinations(cards, n=3):
    "Return all combinations of n of the cards."
    if n==1: 
        return [[card] for card in cards]
    else: 
        return [[cards[i]]+g 
                for i in range(len(cards)-1) for g in combinations(cards[i+1:],n-1)]

def deal(n, deck): 
    "Deal n cards from the deck."
    return [deck.pop() for _ in range(n)]

def play_deal(N=1000000):
    "Simulate N deals, keeping track of Sets and NoSets."
    tallies = defaultdict(int)
    for size in (12, 15):
        for _ in range(N):
            random.shuffle(DECK)
            s = find_set(DECK[:size])
            tallies[size, bool(s)] += 1
    show(tallies, 'first deal')

def play_games(N=50000):
    "Simulate playing N games, keeping track of Sets and NoSets."
    tallies = defaultdict(int)
    for _ in range(N):
        deck = DECK[:]
        random.shuffle(deck)
        table = deal(12, deck)
        while deck or find_set(table):
            s = find_set(table)
            tallies[len(table), bool(s)] += 1
            for card in s: table.remove(card)
            if (s and len(table) < 12 and deck) or (not s): 
                table += deal(3, deck)
    show(tallies, 'playing a game')

def show(tallies, kind):
    "Print out the counts."
    print
    print 'Size |  Set  | NoSet | Set:NoSet for %s' % kind
    print '-----+-------+-------+---------------'
    for size in (12, 15):
        y, n = tallies[size, True], tallies[size, False]
        print '%4d |%6d |%6d | %s:1' % (size, y, n, ('oo' if n==0 else int(round(float(y)/n))))

def test():
    assert len(DECK) == 81 == len(set(DECK))
    assert is_set(['3R=O', '2R=S', '1R=D'])
    assert not is_set(['3R=O', '2R=S', '1R@D'])
    assert find_set(['1POO', '2G=D', '3R=O', '2R=S', '1R=D']) == ['3R=O', '2R=S', '1R=D']
    assert not find_set(['1POO', '2G=D', '3R=O', '2R=S', '1R@D'])
    assert not find_set('2P=O 3P=D 2R=O 3GOO 2POD 3R@D 2ROO 2ROS 1P@S 2P@O 3ROS 2GOD 2P@D 1GOD 3GOS'.split())
    assert combinations([1, 2, 3, 4]) == [[1, 2, 3], [1, 2, 4], [1, 3, 4], [2, 3, 4]]
    print 'All tests pass.'

test()
play_deal()
play_games()