#!/usr/bin/env python
#
# Snow Simulator -- Heavily modified version of the pygame 'stars.py' simple starfield example
# Glen Greer
# 2007

import math
import pygame
import pygame.locals
import random

# Constants
WINSIZE = [640, 480]
NUMFLAKES = 350
SNOW = (255, 255, 255)
BACKGROUND = (20, 20, 40, 255)
ICE = (127, 127, 255, 255)
STATIONARY = (250, 250, 250, 255)
RED = (255, 127, 127)
GREEN = (127, 255, 127)

def init_flake(i):
    return i, SNOW, (random.randrange(WINSIZE[0]), -2 * random.randrange(WINSIZE[1]))

def initialize_flakes():
    return [init_flake(i) for i in range(NUMFLAKES)]

def draw_flakes(surface, flakes, color):
    for i, type, pos in flakes:
        pos = (int(pos[0]), int(pos[1]))
        surface.set_at(pos, color)

def new_pos(surface, pos, dir, y_dir=1):
    new = [pos[0] + dir, pos[1] + y_dir]
    if (new[0] < 0):
        new[0] = WINSIZE[0] - 1
    elif (new[0] >= WINSIZE[0]):
        new[0] = 0
    if (0 <= new[1] < WINSIZE[1]):
        dest = surface.get_at(new)
    else:
        dest = BACKGROUND
    return tuple(new), dest

def move_flakes(surface, flakes):
    for i, type, pos in flakes:
        dir = 1 - (2 * random.randrange(2))
        new, dest = new_pos(surface, pos, dir)
        if (BACKGROUND == dest):
            if (ICE == type):
                up, updest = new_pos(surface, pos, 0, -1)
                if (BACKGROUND == updest):
                    flakes[i] = i, type, new
                else:
                    surface.set_at(pos, ICE)
                    flakes[i] = init_flake(i)
            else:
                flakes[i] = i, type, new
        else:
            type = ICE
            dir *= -1
            new, dest = new_pos(surface, pos, dir)
            if (BACKGROUND == dest):
                flakes[i] = i, type, new
            else:
                dir = 0
                new, dest = new_pos(surface, pos, dir)
                if (BACKGROUND == dest):
                    flakes[i] = i, type, new
                else:
                    surface.set_at(pos, STATIONARY)
                    flakes[i] = init_flake(i)

def main():
    pygame.init()
    random.seed()
    pygame.display.set_caption('Snow')
    screen = pygame.display.set_mode(WINSIZE)

    flakes = initialize_flakes()

    screen.fill(BACKGROUND)
    pygame.draw.line(screen, GREEN, (0, WINSIZE[1] - 1), (WINSIZE[0] - 1, WINSIZE[1] - 1))

    done = 0
    while not done:
        draw_flakes(screen, flakes, BACKGROUND)
        move_flakes(screen, flakes)
        draw_flakes(screen, flakes, SNOW)
        pygame.display.update()
        for e in pygame.event.get():
            if e.type == pygame.locals.QUIT or (e.type == pygame.locals.KEYUP and e.key == pygame.locals.K_ESCAPE):
                done = 1
                break
            elif e.type == pygame.locals.MOUSEBUTTONDOWN and e.button == 1:
                pygame.draw.circle(screen, RED, e.pos, 10)

if __name__ == '__main__':
    main()