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 __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.FULLSCREEN )
		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()