from pygame import *

t = time.Clock()

foo = 3
bar = 2

scr = display.set_mode((500,500))

p = set()
f = True

while f:
    for ev in event.get():
        if ev.type == MOUSEMOTION and ev.buttons[0]:
            p |= set(((ev.pos),))
            scr.set_at((ev.pos),(200,200,200))
            display.update(((ev.pos),(1,1)))
        elif ev.type == KEYDOWN and ev.key == K_RETURN: f = False
        elif ev.type == QUIT:
            quit()
            exit()

g = 0
while True:
    if event.poll().type == QUIT: break
    
    q = {}
    for j,k in p:
        for x,y in((-1,-1),(0,-1),(1,-1),(-1,0),(1,0),(-1,1),(0,1),(1,1)):
            q[(j+x,k+y)] = q.setdefault((j+x,k+y),0)+1
    g += 1
    print('population: %i, generation: %i'%(len(p),g))
    scr.fill(0)
    # *** uncomment the following line for black and white version and triple-quote the colored version ***
    #p = set([scr.set_at(i,(200,200,200))or i for i in q if q[i] == foo or (q[i] == bar and i in p)])
    
    # **** COLORED VERSION ***
    z = set()
    for i in q:
        if q[i] == foo:
            z |= set([i])
            if i in p:
                scr.set_at(i,(250,150,150))
            else:
                scr.set_at(i,(150,150,250))
        elif (q[i] == bar and i in p):
            z |= set([i])
            scr.set_at(i,(250,150,150))
    p = z
    # ***********************
    
    display.flip()
    t.tick(40)
quit()