from itertools import chain , product , repeat , izip , imap
import pygame

def nextgen(generation):
    newgen=set()
    neighbors=lambda (x,y):((x+a,y+b) for a,b in product( [1,0,-1],[1,0,-1] ) if (a,b)!=(0,0) )
    neighborhood=generation|set(chain(*map(neighbors,generation))) # union of all neighbors set and current generation set
    for cell in neighborhood:
        alive=len(set(neighbors(cell)) & generation) # number of alive neighbors of cell: length of the set
        #got by intersection of current generation and neighbors of cell. ie: how many of the neighbors exist(are alive) in the current generation?
        if (alive in [2,3] and cell in generation) or (alive ==3 and cell not in generation):
            newgen.add(cell)
    return newgen

def setup():
    clock=pygame.time.Clock()
    screen=pygame.display.set_mode(resolution)
    return clock,screen

def getdim(generation):
    (maxx,minx),(maxy,miny)=map(lambda p:(max(p),min(p)), zip(*generation))
    generation=set([(x-minx,y-miny) for x,y in generation])   # Normalize
    width,height=maxx-minx+1,maxy-miny+1
    return (width,height,generation)

def getrect((cell,dimension)):
    w,h=dimension
    rwidth=resolution[0]/(30) # dimensions of cells.30 X 30 grid
    rheight=resolution[1]/(30)
    position=(cell[0]*rwidth)+(resolution[0]/2)-((rwidth*w)/2),(cell[1]*rheight)+(resolution[1]/2)-((rheight*h)/2)
    return pygame.Rect(position,(rwidth-2,rheight-2)) #the 2 is padding to distinguish between adjacent cells.

def paint(generation,screen):
    w,h,generation=getdim(generation)
    cells=imap(getrect,izip(generation,repeat( (w,h) ) ) )
    for cell in cells:
        pygame.draw.rect(screen,red,cell,0)
    pygame.display.flip()

red=255,0,0
resolution=(600,600)
clock,screen=setup()
generation=set([(0,0),(1,0),(2,0)]) #blinker
# generation=set([(0,0),(1,0),(2,0),(0,1),(1,1),(2,1),(0,2),(1,2),(2,2)]) # another pattern
while pygame.event.poll().type!=pygame.QUIT:
    clock.tick(500)
    paint(generation,screen)
    generation=nextgen(generation)
    pygame.time.delay(500)
    screen.fill( (0,0,0) )