import pygame
from pygame.locals import *

import Numeric
from Numeric import *

SIZE = array([64, 64])
SCALE = 4
pygame.init()
d = pygame.display.set_mode(SIZE * SCALE) # implies pygame.init()
s = pygame.Surface(SIZE)
r = s.get_rect().inflate(-2, -2)
clock = pygame.time.Clock()
font = pygame.font.SysFont("Courier New", 20)

U8 = UnsignedInt8
particles = []
pixels = zeros(concatenate([SIZE, [3]])) #, UnsignedInt16, savespace=1)
a = pixels[:,:,1] # green plane

rel_avg = array([0.0, 0.0])  # used for smoothing mouse movements

try:
    while True:
        pygame.event.pump()
        # events
        for e in pygame.event.get():
            if e.type in (QUIT, KEYDOWN):
                raise SystemExit()
                
        # smooth the relative speed
        pos = array(pygame.mouse.get_pos())
        rel = array(pygame.mouse.get_rel())
        rel_avg /= 2.0
        rel_avg += rel   # results in rel_avg being bigger than rel
        if sum(pygame.mouse.get_pressed()):
            particles.append((pos / float(SCALE),
                              rel_avg / float(SCALE) / 2))
        
        # update display
        pygame.surfarray.blit_array(s, pixels)
        for p, speed in particles:
            s.set_at(p.astype(Int), (0, 255, 255))
        d.blit(pygame.transform.scale(s, SIZE * SCALE), (0, 0))
        # mouse tail
        tail_pos = pos - SCALE / 2. * rel_avg
        pygame.draw.line(d, (255, 255, 0), tail_pos.astype(Int), pos.astype(Int))

        clock.tick()
        fps = clock.get_fps()
        d.blit(font.render("%.0f ms/f = %.1f f/s" % (1e3 / (fps or 1e6), fps),
                           True, (255, 255, 255)),
               (0, 0))
        pygame.display.flip()
        # move
        particles = [
            (pos + speed, speed)
            for (pos, speed) in particles
            if r.collidepoint(pos + speed) and hypot(*speed) > 0.1]
        # accelerate from field
        particles = [(pos, speed * (1 + 0.0003 * a[pos.astype(Int)]))
                     for (pos, speed) in particles]
        # disturb field
        for pos, speed in particles:
            a[pos.astype(Int)] += 200
        # field blurs with time
        center = a[1:-1, 1:-1]
        neighbours = [a[1:-1, :-2], a[1:-1, 2:], a[:-2, 1:-1], a[2:, 1:-1]]
        blur = sum([6 * center] + neighbours) // 10
        center[:] = clip(blur, 0, 255).astype(U8)

except SystemExit:
    pass
##finally:
pygame.quit()