Source

sketchbook / demo / main.py

Full commit
from __future__ import division
import os
import sys
import random
import math
import cProfile
from optparse import OptionParser

import pygame
from pygame.locals import *
from pygame import time

from plasma import Plasma
from fire import Fire
from rotate import Rotate
from matrix import get_2d_array

APP_NAME = "Old school demo v0.1"
TILE_SIZE = 10
WINDOW_SIZE = (640,480)
FILL_SCREEN = True
PROFILE_ITERATIONS = 100

FONT_SIZE = 16
###################################################################
# Problems
# Black line on plasma
# Rotate doesn't get back to flat
###################################################################

class DemoController(object):
    def __init__(self, profile=False):
        """If profile is set, the simulation will run for 100 iterations and quit displaying cProfile output"""
        #pygame
        if FILL_SCREEN:
            self._screen = pygame.display.set_mode((0,0),FULLSCREEN)
        else:
            self._screen = pygame.display.set_mode((WINDOW_SIZE[0],WINDOW_SIZE[1]), FILL_SCREEN)
        self._font = pygame.font.SysFont("DejaVu Sans Mono", FONT_SIZE)
        
        self.bg = pygame.display.get_surface().convert()

        screen_size = self._screen.get_size()
        self.grid_size = (int(screen_size[0]/TILE_SIZE), int(screen_size[1]/TILE_SIZE))
        self.grid = get_2d_array(*self.grid_size)

        self.paused = False
        self.clock = time.Clock()
        self.profile = profile
        self.fps = 0

        #cache spectrum colours, looking up on the fly is slow
        self.fire_spectrum = []
        spectrum = pygame.image.load('fire.bmp')
        for i in xrange(0,255):
            c = spectrum.get_at( (i, 1) )
            h = pygame.Color(*c)
            self.fire_spectrum.append(h)

        # (Class, Iterations, Transition)
        self.effects = [
            (Plasma,300, {}),
            (Rotate, 180, {'tile_size':TILE_SIZE}),
            (Fire, 150, {}),
            ]
        self.iteration = 0

        self.effect_index = 0
        self.effect_object = self.get_effect()

        self.render()

    def get_effect(self):
        effect = self.effects[self.effect_index]
        return effect[0](self.grid, effect[1], **effect[2])

    def update(self):
        if not self.effect_object:
            self.effect_object = self.get_effect()
        self.grid = self.effect_object.get_grid()
        self.render()

        if self.effect_object.iteration < 1:
            self.effect_index += 1
            if self.effect_index > len(self.effects)-1:
                self.effect_index = 0
            self.effect_object = self.get_effect()

    def render(self):
        spec = self.fire_spectrum
        grid = self.grid
        surf = self.bg
        rect = pygame.Rect(0,0,TILE_SIZE,TILE_SIZE)

        for y, row in enumerate(grid):
            rect.top = y*TILE_SIZE
            for x, val in enumerate(row):
                rect.left = x * TILE_SIZE
                self._screen.fill(spec[val], rect)
        #self._screen.blit(self.bg, (0,0))
        pygame.display.flip()
        self.iteration += 1

    def input(self,events):
        """Deals with kepxesses and mouse events"""
        for event in events:
            if hasattr(event, 'unicode') and event.unicode == 'q':
                sys.exit(0)

            if hasattr(event, 'unicode') and event.unicode == 'p':
                self.paused = not self.paused
                return

            if hasattr(event, 'unicode') and event.unicode == 'n':
                if self.paused:
                    self.update()

    def mainloop(self):
        """Main loop!"""
        done = False
        while not done:
            self.input(pygame.event.get())
            if not self.paused:
                self.clock.tick()
                self.fps = self.clock.get_fps()            
                self.update()
            if self.profile and self.iteration > PROFILE_ITERATIONS:
                done = True

if __name__ == "__main__":
    parser = OptionParser()
    parser.add_option("--profile", action="store_true", dest="profile", default="False", help="Profile 100 iterations (default: False)")
    (options, args) = parser.parse_args()

    pygame.init()
    pygame.display.set_caption(APP_NAME)
    pc = DemoController()

    if options.profile is True:
        pc.profile = True
        cProfile.run('pc.mainloop()')
        pygame.quit()
        raw_input("Press Enter to exit")        
    else:
        pc.mainloop()
        pygame.quit()