one-file-projects/colortest.py

import pygame
from pygame.locals import *
from random import random,randint
from math import sqrt

screen = None
color = (128,128,128)
color_vel = (0,0,0)
running = True

class Vector(tuple):
        def __add__(self,other):
                if len(self) != len(other):
                        raise Exception("Can't add vectors of different size: %d %d" % (len(self), len(other)))
                return Vector([ self[i]+other[i] for i in range(len(self)) ])

        def __sub__(self,other):
                if len(self) != len(other):
                        raise Exception("Can't subtract vectors of different size")
                return Vector([ self[i]-other[i] for i in range(len(self)) ])

        def __mul__(self,factor):
                return Vector([ self[i]*factor for i in range(len(self)) ])

        def __div__(self,factor):
                return self.__mul__(1.0/factor)

        def __truediv__(self,factor):
                return self.__div__(factor)

        def __rmul__(self,factor):
                return self.__mul__(factor)

        def __rdiv__(self,factor):
                raise NotImplemented

        def __iadd__(self,other):
                return self.__add__(other)

        def __isub__(self,other):
                return self.__sub__(other)

        def __neg__(self):
                return self.__mul__(-1.0)

        def __setitem__(self,k,v):
                raise Exception("Not available") # immutable

        def distance(self, other):
                if len(self) != len(other):
                        raise Exception("Can't calculate distance between vectors of different size")
                return sqrt(sum([ (self[i]-other[i])**2 for i in range(len(self)) ]))

        def length(self):
                return self.distance( Vector( (0,0) ) )


class Tick(object):
        def __init__(self, tickrange):
                self.tickrange = tickrange
                self.m_tick = 0

        def tick(self):
                self.m_tick = (self.m_tick+1)%self.tickrange
                return self.m_tick == 0

class Mode(object):
        def __init__(self, screensize):
                self.screensize = screensize

        def update(self):
                pass

        def draw(self,screen):
                pass

def minmax(vmin,vmax):
        return lambda v: min(vmax,max(vmin,v))

class TransitionMode(Mode):
        t = 0.3
        a = 30
        def __init__(self, screensize):
                Mode.__init__(self,screensize)
                self.color = Vector( (127,127,127) )
                self.color_vel = Vector( 2*self.a*random()+self.a for i in range(3) )
                self.utick = Tick(10)

        def update(self):
                if self.utick.tick():
                        randvec = Vector( (2*self.a*random()-self.a) for i in range(3) )
                        self.color_vel = (1-self.t)*self.color_vel+self.t*randvec
                self.color = Vector( map(minmax(0,255), self.color+self.color_vel) )

        def draw(self,screen):
                screen.fill( self.color )

class LelMode(Mode):
        color = [(255,0,0),(0,255,0)]
        size = 10
        def __init__(self,screensize):
                Mode.__init__(self,screensize)
                self.offset = 0
                self.tick = 0

        def update(self):
                self.offset = (self.offset+1)%2

        def draw(self,screen):
                w,h = screen.get_size()
                for x in range(w):
                        color = self.color[ (x//self.size + self.offset)%2]
                        pygame.draw.line(screen, color, (x,0), (x,h) )

class BallMode(Mode):
        def __init__(self,screensize):
                Mode.__init__(self,screensize)
                self.ball = [200,200]
                self.radius = 70
                self.velocity = [10,10]

        def update(self):
                w, h = self.screensize
                if self.ball[0] + self.radius >= w:
                        self.velocity[0] = -abs(self.velocity[0])
                elif self.ball[0] - self.radius < 0:
                        self.velocity[0] = abs(self.velocity[0])
                elif self.ball[1] + self.radius >= h:
                        self.velocity[1] = -abs(self.velocity[1])
                elif self.ball[1] - self.radius < 0:
                        self.velocity[1] = abs(self.velocity[1])

                for i in range(2):
                        self.ball[i] += self.velocity[i]

        def draw(self,screen):
                screen.fill( (0,0,0) )
                pygame.draw.circle(screen, (255,0,0), tuple(self.ball), self.radius)


class Particle:
        def __init__(self,color, pos, vel):
                self.color = color
                self.pos = Vector(pos)
                self.vel = Vector(vel)

        def updateVel(self,gravity,max_v=20):
                gravity = Vector(gravity)

                #Update velocity
                self.vel = 0.95*self.vel+0.01*random()*(gravity-self.pos)+0.1*random()*Vector( randint(-70,70) for i in range(2) )

                #Limit speed
                l = self.vel.length()
                self.vel *= max_v/l

        def updatePos(self):
                self.pos = self.pos + self.vel

        def draw(self,screen):
                w,h = 1280,800
                if self.pos[0] > 0 and self.pos[0] < w and self.pos[1] > 0 and self.pos[1] < h:
                        pos = tuple(map(int,self.pos))
                        velpos = tuple(map(int,self.pos+self.vel))
                        pygame.draw.circle(screen, self.color, pos, 5)
                        pygame.draw.line(screen, (255,255,255), pos, velpos, 1 )

        @staticmethod
        def generate(xmax,ymax):
                return Particle( tuple( randint(100,255) for i in range(3) ), ( randint(0,xmax), randint(0,ymax) ), tuple( randint(-50,50) for i in range(2) ) )

class SwarmMode(Mode):
        swarmsize = 60
        def __init__(self,screensize):
                Mode.__init__(self,screensize)
                self.particles = list( Particle.generate(1280,800) for i in range(self.swarmsize) )
                self.generateGravity()
                self.gtick = Tick(150)
                self.utick = Tick(2)

        def generateGravity(self):
                #self.gravity = tuple( randint(0,m) for m in [1280,800] )
                self.gravity = Vector( (1280,800) )/2

        def update(self):
                if self.gtick.tick():
                        self.generateGravity()

                v = self.utick.tick()

                #king = self.particles[0]
                #if v: king.updateVel(self.gravity, max_v=3)
                #king.updatePos()

                for p in self.particles:
                        if v: p.updateVel(self.gravity)
                        p.updatePos()

        def draw(self,screen):
                screen.fill( (0,0,0) )
                for p in self.particles:
                        p.draw(screen)

class Game(object):
        def __init__(self, modes=[SwarmMode,TransitionMode,LelMode,BallMode]):
                self.modes = modes
                pygame.init()
                self.screen = pygame.display.set_mode( (1280,800), pygame.HWSURFACE )
                pygame.mouse.set_visible(False)
                self.setmode(0)
                self.running = True

        def update(self):
                self.mode.update()

        def setmode(self,i):
                self.modeindex = i
                self.mode = self.modes[i%len(self.modes)](self.screen.get_size())

        def nextmode(self):
                self.setmode( self.modeindex+1 )

        def draw(self):
                self.mode.draw(self.screen)
                pygame.display.flip()

        def getevents(self):
                for event in pygame.event.get():
                        if event.type == QUIT:
                                self.running = False
                        elif event.type == KEYDOWN:
                                if event.key == K_ESCAPE:
                                        pygame.event.post(pygame.event.Event(QUIT))
                                elif event.key == K_SPACE:
                                        self.nextmode()

        def run(self):
                clock = pygame.time.Clock()
                while self.running:
                        clock.tick(30)
                        self.getevents()
                        self.update()
                        self.draw()


if __name__ == '__main__':
        Game().run()