#!/usr/bin/env python  
# -*- coding: utf-8 -*-

#=========================================================
# polen.py - N partículas de igual masa chocando 
# con una de mayor masa (movimiento browniano)
#---------------------------------------------------------
# FJA - neocipres@gmail.com    Mayo de 2010
#=========================================================
 
import sys, random
import pygame
from pygame.locals import *

XMAX  = 800 
YMAX = 600
SIZE = XMAX ,  YMAX
NEGRO = (0, 0, 0)
M1 = 5
M2 = 1

class Polen(pygame.sprite.Sprite):
    def __init__(self):
        pygame.sprite.Sprite.__init__(self)
        self.image = pygame.image.load( "/home/paco/prgkgnu/python/pygame/polen.png" )
        self.rect = self.image.get_rect()
        self.rect.center = XMAX/2,  YMAX/2
        self.vx = 0
        self.vy = 0        
        
    def update(self,  grupo):
        
        self.rect.move_ip((self.vx, self.vy))
        
        if self.rect.left < 0 or self.rect.right > XMAX:
                self.vx = -self.vx
                self.vy = self.vy
        if self.rect.top < 0 or self.rect.bottom > YMAX:
                self.vx = self.vx
                self.vy = -self.vy      
                
        choques = pygame.sprite.spritecollide(self, grupo,  0)  # 
        for choque in choques:
           
            if choque != self:
                u1x = self.vx       # Principio de conservación de la cantidad de movimiento 
                u1y = self.vy
                u2x = choque.vx
                u2y = choque.vy
                vcx = float((M1*u1x + M2*u2x)/(M1 + M2))
                vcy = float((M1*u1y + M2*u2y)/(M1 + M2))
                choque.vx = round(2*vcx - u2x)
                choque.vy = round(2*vcy - u2y)
                if choque.vx == 0 and choque.vy == 0: choque.kill()
                self.vx =  round(2*vcx - u1x)
                self.vy =  round(2*vcy - u1y)
                
                
                
class Bola(pygame.sprite.Sprite):
    def __init__(self):
        pygame.sprite.Sprite.__init__(self)
        self.image = pygame.image.load( "/home/paco/prgkgnu/python/pygame/bola.gif" )
        self.rect = self.image.get_rect()
        self.rect.center = random.randrange(XMAX  -  50) + 50, random.randrange(YMAX - 50) + 50
        self.vx = random.choice([-2, -1,  1,  2])
        self.vy = random.choice([2, 1,  0, -1,  -2])
        
        
    def update(self,  grupo):
        
        self.rect.move_ip((self.vx, self.vy))
        
        if self.rect.left < 0 or self.rect.right > XMAX:
                self.vx = -self.vx
                self.vy = self.vy
        if self.rect.top < 0 or self.rect.bottom > YMAX:
                self.vx = self.vx
                self.vy = -self.vy      
                
        choques = pygame.sprite.spritecollide(self, grupo,  0)  # Para detectar las colisiones entre los sprites (bolas)
        for choque in choques:
           
            if choque != self:
                u1x = self.vx       # Principio de conservación de la cantidad de movimiento para dos partículas de igual masa
                u1y = self.vy
                u2x = choque.vx
                u2y = choque.vy
                self.vx = u2x
                self.vy = u2y
                choque.vx = u1x
                choque.vy = u1y
                if choque.rect.left < 0 or choque.rect.right > XMAX:
                    choque.kill()
                elif choque.rect.top < 0 or choque.rect.bottom > YMAX:
                    choque.kill()
               
          
                
        
        
def main():
    run = True
    random.seed()
    pygame.init()
    screen = pygame.display.set_mode( (XMAX, YMAX) )
    screen.fill(NEGRO)
    pygame.display.set_caption('movimiento browniano')
    bolaSprites = pygame.sprite.RenderClear()       # Para crear el contenedor de los sprites
    polenSprites = pygame.sprite.RenderClear()
    
    polen = Polen()
    polenSprites.add(polen)
    
    try:
        cadnum = sys.argv[1]
        n = int(cadnum)
        if n > 100:
            n = 100
    except IndexError:
        n = 50
    for n in range(n):
        bola = Bola()
        bolaSprites.add(bola)
        
    while run:
        for e in pygame.event.get():
            if (e.type == QUIT):
                run = False
                break
            elif (e.type == KEYDOWN):
                if( e.key == K_ESCAPE or e.key == K_q):
                    run = False
                    break
                if( e.key == K_f ):
                    pygame.display.toggle_fullscreen()
        
        
            
        bolaSprites.update(bolaSprites)                                # Actualizar los sprites
        polenSprites.update(bolaSprites)
        
           
        bolaSprites.clear(screen, pygame.Surface(SIZE))     # Limpia la pantalla y dibuja de nuevo 
        bolaSprites.draw(screen)
        polenSprites.clear(screen, pygame.Surface(SIZE)) 
        polenSprites.draw(screen)
            
            
        pygame.display.flip()
        
        
if __name__ == '__main__':
    main()