from pygame import *
t = time.Clock()
font.init()
font = font.Font(None,20)

scr = display.set_mode((500,500))
scr.fill(0)
display.flip()

def add_cell():
    global p
    x,y = mouse.get_pos()
    x = (x-offx)//zoom
    y = (y-offy)//zoom
    p |= set(((x,y),))
    display.update(scr.fill((150,150,250),(x*zoom+offx,y*zoom+offy,zoom,zoom)))
    if pause: display.update(scr.blit(font.render('population: %i '%len(p),1,(200,200,200),(50,50,100)),(0,0)))

def redraw():
    h = Q[(g&1)^1]
    scr.fill(0)
    [scr.fill((150,150,250),(i[0]*zoom+offx,i[1]*zoom+offy,zoom,zoom)) for i in p.difference(p2)]
    [scr.fill((150,150,250)if h[i] == 3 and not i in p else (250,150,150),(i[0]*zoom+offx,i[1]*zoom+offy,zoom,zoom)) for i in h if h[i] == 3 or (i in p and h[i] == 2)]
    scr.blit(font.render('population: %i '%len(p),1,(200,200,200),(50,50,100)),(0,0))
    scr.blit(font.render('generation: %i '%g,1,(200,200,200),(50,50,100)),(0,20))
    scr.blit(font.render('elapsed time: %i'%elaps,1,(200,200,200),(50,50,100)),(0,40))
    scr.blit(font.render('FPS: %i '%t.get_fps(),1,(200,200,200),(50,50,100)),(0,60))
    display.flip()
    
zoom = 1
offx,offy = 0,0
pause = True
elaps = 0

p = frozenset()
#with open('cellules.txt') as p: p = eval(p.read())
p2 = frozenset()
Q = [{},{}]
g = 0
redraw()

while zoom:
    red = 0
    for ev in event.get():
        if ev.type == MOUSEMOTION and ev.buttons[0]:
            if  key.get_pressed()[K_SPACE]:
                add_cell()
            else:
                x,y = ev.rel
                offx += x
                offy += y
                if pause: red = 1
        elif ev.type == MOUSEBUTTONDOWN:
            if ev.button == 1 and key.get_pressed()[K_SPACE]:
                add_cell()
                continue
            z = 0
            if ev.button == 4: z = 1
            elif ev.button == 5 and zoom>1: z = -1
            x,y = ev.pos
            x0 = (x-offx)//zoom
            y0 = (y-offy)//zoom
            zoom += 1 * z
            offx = x-x0*zoom
            offy = y-y0*zoom
            if pause: red = 1
        elif ev.type == QUIT:
            zoom = 0
            break
        elif ev.type == KEYDOWN:
            if ev.key == K_ESCAPE:
                quit()
                exit()
            elif ev.key == K_p:
                pause = not pause
                if pause:
                    p2 = p.copy()
                    redraw()
                else:
                    t.tick()
                    #with open('cellules.txt','w') as pp: pp.write(str(p))
    else:
        if pause:
            if red : redraw()
            continue
        g += 1
        Q[g&1] = {}
        q = Q[g&1]
        for j,k in p:
            
            x = j-1,k-1
            q[x] = q.setdefault(x,0)+1
            x = j,k-1
            q[x] = q.setdefault(x,0)+1
            x = j+1,k-1
            q[x] = q.setdefault(x,0)+1
            x = j-1,k
            q[x] = q.setdefault(x,0)+1
            x = j+1,k
            q[x] = q.setdefault(x,0)+1
            x = j-1,k+1
            q[x] = q.setdefault(x,0)+1
            x = j,k+1
            q[x] = q.setdefault(x,0)+1
            x = j+1,k+1
            q[x] = q.setdefault(x,0)+1
        
    elaps += t.tick(40)
    scr.fill(0)
    scr.blit(font.render('population: %i'%len(p),1,(200,200,200)),(0,0))
    scr.blit(font.render('generation: %i'%g,1,(200,200,200)),(0,20))
    scr.blit(font.render('elapsed time: %i'%elaps,1,(200,200,200)),(0,40))
    scr.blit(font.render('FPS: %i'%t.get_fps(),1,(200,200,200)),(0,60))
    p = frozenset([scr.fill((150,150,250)if j == 3 and not (x,y) in p else (250,150,150),(x*zoom+offx,y*zoom+offy,zoom,zoom))==0 or (x,y) for (x,y),j in q.items() if j == 3 or ((x,y) in p and j == 2)])
    display.flip()
#with open('cellulesoutput.txt','w') as pp: pp.write(str(p))   
quit()
