#!/usr/bin/env python
import numpy as np
from itertools import product, count
from threading import Thread, Lock
import os
import sys
from chroma.color import map_wavelength, map_to_color
from chroma.geometry import Mesh, Solid, Geometry
import chroma.make as make
import matplotlib.pyplot as plt

from chroma.gpu import GPU, CUDAFuncs, to_float3
from chroma.tools import timeit
from chroma.transform import rotate
from chroma.optics import vacuum, lambertian_surface
import chroma.project as project

import pygame
from pygame.locals import *

from pycuda import gpuarray as ga
from pycuda.characterize import sizeof
import pycuda.driver as cuda

from subprocess import call
import shutil
import tempfile

def buildable(identifier):
    """
    Create a decorator which tags a function as buildable and assigns the
    identifying string `identifier`.

    Example:
        >>> @buildable('my_sphere')
        >>> def build_my_sphere():
        >>>     g = Geometry()
        >>>     g.add_solid(Solid(sphere(), vacuum, water))
        >>>     return g
    """
    def tag_as_buildable(func):
        func.buildable = True
        func.identifier = identifier
        return func
    return tag_as_buildable

def build(obj, bits):
    """Construct and build a geometry from `obj`."""
    if inspect.isfunction(obj):
        try:
            if obj.buildable:
                obj = obj()
        except AttributeError:
            raise Exception('function %s is not buildable.' % obj.__name__)

    if isinstance(obj, Geometry):
        geometry = obj
    elif isinstance(obj, Solid):
        geometry = Geometry()
        geometry.add_solid(obj)
    elif isinstance(obj, Mesh):
        geometry = Geometry()
        geometry.add_solid(Solid(obj, vacuum, vacuum, surface=lambertian_surface, color=0x99ffffff))
    else:
        raise Exception('cannot build type %s' % type(obj))

    geometry.build(bits)

    return geometry

def bvh_mesh(geometry, layer):
    lower_bounds = geometry.lower_bounds[geometry.layers == layer]
    upper_bounds = geometry.upper_bounds[geometry.layers == layer]

    if len(lower_bounds) == 0 or len(upper_bounds) == 0:
        raise Exception('no nodes at layer %i' % layer)

    dx, dy, dz = upper_bounds[0] - lower_bounds[0]
    center = np.mean([upper_bounds[0],lower_bounds[0]], axis=0)

    mesh = make.box(dx, dy, dz, center)

    for center, dx, dy, dz in zip(np.mean([lower_bounds,upper_bounds],axis=0),*zip(*upper_bounds-lower_bounds))[1:]:
        mesh += make.box(dx,dy,dz,center)

    return mesh

def encode_movie(dir):
    root, ext = 'movie', 'avi'
    for i in count():
        filename = '.'.join([root + str(i).zfill(5), ext])

        if not os.path.exists(filename):
            break

    call(['mencoder', 'mf://' + dir + '/*.png', '-mf', 'fps=10', '-o', filename, '-ovc', 'xvid', '-xvidencopts', 'bitrate=3000'])

    shutil.rmtree(dir)

    print 'movie saved to %s.' % filename

def get_rays(position, size = (800, 600), width = 0.035, focal_length=0.018):
    height = width*(size[1]/float(size[0]))

    x = np.linspace(-width/2, width/2, size[0])
    z = np.linspace(-height/2, height/2, size[1])

    grid = np.array(tuple(product(x,[0],z)))

    grid += (0,focal_length,0)
    focal_point = np.zeros(3)

    grid += position
    focal_point += position

    return grid, focal_point-grid

class Camera(Thread):
    def __init__(self, geometry, size=(800,600), device_id=None, enable3d=False):
        Thread.__init__(self)
        self.geometry = geometry
        self.device_id = device_id
        self.size = size
        self.enable3d = enable3d

        self.unique_bvh_layers = np.unique(self.geometry.layers)
        self.currentlayer = None
        self.bvh_layers = {self.currentlayer : self.geometry}

        try: 
            import spnav as spnav_module
            self.spnav_module = spnav_module
            self.spnav = True
        except:
            self.spnav = False

    def init_gpu(self):
        self.gpu = GPU(self.device_id)
        self.gpu.load_geometry(geometry)

        self.width, self.height = size

        pygame.init()
        self.screen = pygame.display.set_mode(size)
        pygame.display.set_caption('')
        self.clock = pygame.time.Clock()

        self.doom_mode = False
        try:
            if self.width == 640: # SECRET DOOM MODE!
                print 'shotgun activated!'
                self.doom_hud = pygame.image.load('images/doomhud.png').convert_alpha()
                rect = self.doom_hud.get_rect()
                self.doom_rect = rect.move(0, self.height - rect.height)
                self.doom_mode = True
        except:
            pass

        lower_bound, upper_bound = self.geometry.mesh.get_bounds()
        self.scale = np.linalg.norm(upper_bound-lower_bound)

        self.nblocks = 64

        self.point = np.array([0, -self.scale*1.0, (lower_bound[2]+upper_bound[2])/2])
        self.axis1 = np.array([0,0,1], float)
        self.axis2 = np.array([1,0,0], float)

        self.film_width = 0.035

        if self.enable3d:
            self.point1 = self.point-(self.scale/60,0,0)
            self.point2 = self.point+(self.scale/60,0,0)

            self.viewing_angle = 0.0

            origins1, directions1 = project.from_film(self.point1, size=self.size, width=self.film_width)
            origins2, directions2 = project.from_film(self.point2, size=self.size, width=self.film_width)

            origins = np.concatenate([origins1,origins2])
            directions = np.concatenate([directions1,directions2])

            scope_origins, scope_directions = project.from_film(self.point, size=np.array(self.size)/4.0, width=self.film_width/4.0)

            self.scope_origins_gpu = ga.to_gpu(to_float3(scope_origins))
            self.scope_directions_gpu = ga.to_gpu(to_float3(scope_directions))

            self.distances_gpu = ga.empty(self.scope_origins_gpu.size, dtype=np.float32)
        else:
            origins, directions = project.from_film(self.point, size=self.size, width=self.film_width)

        self.origins_gpu = ga.to_gpu(to_float3(origins))
        self.directions_gpu = ga.to_gpu(to_float3(directions))
        self.pixels_gpu = ga.zeros(self.origins_gpu.size, dtype=np.int32)

        self.alpha = True
        self.movie = False
        self.movie_index = 0
        self.movie_dir = None
        self.render = False

    @timeit
    def initialize_render(self):
        self.rng_states_gpu = cuda.mem_alloc(self.origins_gpu.size*sizeof('curandStateXORWOW', '#include <curand_kernel.h>'))
        self.gpu.kernels.init_rng(np.int32(self.origins_gpu.size), self.rng_states_gpu, np.int32(0), np.int32(0), block=(self.nblocks,1,1), grid=(self.origins_gpu.size//self.nblocks+1,1))
        self.xyz_lookup1_gpu = ga.zeros(len(self.geometry.mesh.triangles), dtype=ga.vec.float3)
        self.xyz_lookup2_gpu = ga.zeros(len(self.geometry.mesh.triangles), dtype=ga.vec.float3)
        self.image_gpu = ga.zeros(self.pixels_gpu.size, dtype=ga.vec.float3)
        self.gpu.context.synchronize()

        self.source_position = self.point

        self.nimages = 0
        self.nlookup_calls = 0
        self.max_steps = 10

    def clear_xyz_lookup(self):
        self.xyz_lookup1_gpu.fill(ga.vec.make_float3(0.0,0.0,0.0))
        self.xyz_lookup2_gpu.fill(ga.vec.make_float3(0.0,0.0,0.0))

        self.nlookup_calls = 0

    def update_xyz_lookup(self, source_position):
        for i in range(self.xyz_lookup1_gpu.size//(self.origins_gpu.size)+1):
            self.gpu.kernels.update_xyz_lookup(np.int32(self.origins_gpu.size), np.int32(self.xyz_lookup1_gpu.size), np.int32(i*self.origins_gpu.size), ga.vec.make_float3(*source_position), self.rng_states_gpu, np.float32(685.0), ga.vec.make_float3(1.0,0.0,0.0), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.origins_gpu.size//self.nblocks+1,1))

        for i in range(self.xyz_lookup1_gpu.size//(self.origins_gpu.size)+1):
            self.gpu.kernels.update_xyz_lookup(np.int32(self.origins_gpu.size), np.int32(self.xyz_lookup1_gpu.size), np.int32(i*self.origins_gpu.size), ga.vec.make_float3(*source_position), self.rng_states_gpu, np.float32(545.0), ga.vec.make_float3(0.0,1.0,0.0), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.origins_gpu.size//self.nblocks+1,1))

        for i in range(self.xyz_lookup1_gpu.size//(self.origins_gpu.size)+1):
            self.gpu.kernels.update_xyz_lookup(np.int32(self.origins_gpu.size), np.int32(self.xyz_lookup1_gpu.size), np.int32(i*self.origins_gpu.size), ga.vec.make_float3(*source_position), self.rng_states_gpu, np.float32(445.0), ga.vec.make_float3(0.0,0.0,1.0), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.origins_gpu.size//self.nblocks+1,1))

        self.nlookup_calls += 1

    def clear_image(self):
        self.image_gpu.fill(ga.vec.make_float3(0.0,0.0,0.0))

        self.nimages = 0

    def update_image(self):
        self.gpu.kernels.update_xyz_image(np.int32(self.origins_gpu.size), self.rng_states_gpu, self.origins_gpu, self.directions_gpu, np.float32(685.0), ga.vec.make_float3(1.0,0.0,0.0), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, self.image_gpu, np.int32(self.nlookup_calls), np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.origins_gpu.size//self.nblocks+1,1))

        self.gpu.kernels.update_xyz_image(np.int32(self.origins_gpu.size), self.rng_states_gpu, self.origins_gpu, self.directions_gpu, np.float32(545.0), ga.vec.make_float3(0.0,1.0,0.0), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, self.image_gpu, np.int32(self.nlookup_calls), np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.origins_gpu.size//self.nblocks+1,1))

        self.gpu.kernels.update_xyz_image(np.int32(self.origins_gpu.size), self.rng_states_gpu, self.origins_gpu, self.directions_gpu, np.float32(445.0), ga.vec.make_float3(0.0,0.0,1.0), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, self.image_gpu, np.int32(self.nlookup_calls), np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.origins_gpu.size//self.nblocks+1,1))

        self.nimages += 1

    def process_image(self):
        self.gpu.kernels.process_image(np.int32(self.pixels_gpu.size), self.image_gpu, self.pixels_gpu, np.int32(self.nimages), block=(self.nblocks,1,1), grid=((self.pixels_gpu.size)//self.nblocks+1,1))

    def screenshot(self, dir='', start=0):
        root, ext = 'screenshot', 'png'

        for i in count(start):
            filename = os.path.join(dir, '.'.join([root + str(i).zfill(5), ext]))

            if not os.path.exists(filename):
                break

        pygame.image.save(self.screen, filename)
        print 'image saved to %s' % filename

    def rotate(self, phi, n):
        self.gpu.kernels.rotate(np.int32(self.pixels_gpu.size), self.origins_gpu, np.float32(phi), ga.vec.make_float3(*n), block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks+1,1))
        self.gpu.kernels.rotate(np.int32(self.pixels_gpu.size), self.directions_gpu, np.float32(phi), ga.vec.make_float3(*n), block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks+1,1))

        self.point = rotate(self.point, phi, n)
        self.axis1 = rotate(self.axis1, phi, n)
        self.axis2 = rotate(self.axis2, phi, n)

        if self.enable3d:
            self.gpu.kernels.rotate(np.int32(self.scope_origins_gpu.size), self.scope_origins_gpu, np.float32(phi), ga.vec.make_float3(*n), block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks+1,1))
            self.gpu.kernels.rotate(np.int32(self.scope_directions_gpu.size), self.scope_directions_gpu, np.float32(phi), ga.vec.make_float3(*n), block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks+1,1))

            self.point1 = rotate(self.point1, phi, n)
            self.point2 = rotate(self.point2, phi, n)

        if self.render:
            self.clear_image()

        self.update()

    def rotate_around_point(self, phi, n, point, redraw=True):
        self.gpu.kernels.rotate_around_point(np.int32(self.origins_gpu.size), self.origins_gpu, np.float32(phi), ga.vec.make_float3(*n), ga.vec.make_float3(*point), block=(self.nblocks,1,1), grid=(self.origins_gpu.size//self.nblocks+1,1))
        self.gpu.kernels.rotate(np.int32(self.directions_gpu.size), self.directions_gpu, np.float32(phi), ga.vec.make_float3(*n), block=(self.nblocks,1,1), grid=(self.directions_gpu.size//self.nblocks+1,1))

        self.axis1 = rotate(self.axis1, phi, n)
        self.axis2 = rotate(self.axis2, phi, n)

        if self.enable3d:
            self.gpu.kernels.rotate_around_point(np.int32(self.scope_origins_gpu.size), self.scope_origins_gpu, np.float32(phi), ga.vec.make_float3(*n), ga.vec.make_float3(*point), block=(self.nblocks,1,1), grid=(self.origins_gpu.size//self.nblocks+1,1))
            self.gpu.kernels.rotate(np.int32(self.scope_directions_gpu.size), self.scope_directions_gpu, np.float32(phi), ga.vec.make_float3(*n), block=(self.nblocks,1,1), grid=(self.directions_gpu.size//self.nblocks+1,1))

        if redraw:
            if self.render:
                self.clear_image()

            self.update()

    def translate(self, v, redraw=True):
        self.gpu.kernels.translate(np.int32(self.pixels_gpu.size), self.origins_gpu, ga.vec.make_float3(*v), block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks,1))

        self.point += v

        if self.enable3d:
            self.gpu.kernels.translate(np.int32(self.scope_origins_gpu.size), self.scope_origins_gpu, ga.vec.make_float3(*v), block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks,1))

            self.point1 += v
            self.point2 += v

        if redraw:
            if self.render:
                self.clear_image()

            self.update()

    def update_pixels(self):
        if self.render:
            while self.nlookup_calls < 10:
                self.update_xyz_lookup(self.source_position)
            self.update_image()
            self.process_image()
        else:
            if self.alpha:
                self.gpu.kernels.ray_trace_alpha(np.int32(self.pixels_gpu.size), self.origins_gpu, self.directions_gpu, self.pixels_gpu, block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks+1,1))
            else:
                self.gpu.kernels.ray_trace(np.int32(self.pixels_gpu.size), self.origins_gpu, self.directions_gpu, self.pixels_gpu, block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks+1,1))

    def update(self):
        if self.enable3d:
            self.gpu.kernels.distance_to_mesh(np.int32(self.scope_origins_gpu.size), self.scope_origins_gpu, self.scope_directions_gpu, self.distances_gpu, block=(self.nblocks,1,1), grid=(self.scope_origins_gpu.size//self.nblocks,1))

            baseline = ga.min(self.distances_gpu).get().item()

            if baseline < 1e9:
                d1 = self.point1 - self.point
                v1 = d1/np.linalg.norm(d1)
                v1 *= baseline/60 - np.linalg.norm(d1)

                self.gpu.kernels.translate(np.int32(self.pixels_gpu.size//2), self.origins_gpu[:self.pixels_gpu.size//2], ga.vec.make_float3(*v1), block=(self.nblocks,1,1), grid=((self.pixels_gpu.size//2)//self.nblocks+1,1))

                self.point1 += v1

                d2 = self.point2 - self.point
                v2 = d2/np.linalg.norm(d2)
                v2 *= baseline/60 - np.linalg.norm(d2)

                self.gpu.kernels.translate(np.int32(self.pixels_gpu.size//2), self.origins_gpu[self.pixels_gpu.size//2:], ga.vec.make_float3(*v2), block=(self.nblocks,1,1), grid=((self.pixels_gpu.size//2)//self.nblocks+1,1))

                self.point2 += v2

            direction = np.cross(self.axis1,self.axis2)
            direction /= np.linalg.norm(direction)
            direction1 = self.point + direction*baseline - self.point1
            direction1 /= np.linalg.norm(direction1)

            new_viewing_angle = np.arccos(direction1.dot(direction))

            phi = new_viewing_angle - self.viewing_angle

            self.gpu.kernels.rotate_around_point(np.int32(self.pixels_gpu.size//2), self.origins_gpu[:self.pixels_gpu.size//2], np.float32(phi), ga.vec.make_float3(*self.axis1), ga.vec.make_float3(*self.point1), block=(self.nblocks,1,1), grid=((self.pixels_gpu.size//2)//self.nblocks+1,1))
            self.gpu.kernels.rotate_around_point(np.int32(self.pixels_gpu.size//2), self.origins_gpu[self.pixels_gpu.size//2:], np.float32(-phi), ga.vec.make_float3(*self.axis1), ga.vec.make_float3(*self.point2), block=(self.nblocks,1,1), grid=((self.pixels_gpu.size//2)//self.nblocks+1,1))
            self.gpu.kernels.rotate(np.int32(self.pixels_gpu.size//2), self.directions_gpu[:self.pixels_gpu.size//2], np.float32(phi), ga.vec.make_float3(*self.axis1), block=(self.nblocks,1,1), grid=((self.pixels_gpu.size//2)//self.nblocks+1,1))
            self.gpu.kernels.rotate(np.int32(self.pixels_gpu.size//2), self.directions_gpu[self.pixels_gpu.size//2:], np.float32(-phi), ga.vec.make_float3(*self.axis1), block=(self.nblocks,1,1), grid=((self.pixels_gpu.size//2)//self.nblocks+1,1))

            self.viewing_angle = new_viewing_angle

        self.update_pixels()

        pixels = self.pixels_gpu.get()

        if self.enable3d:
            pixels1 = pixels[:len(pixels)//2]
            pixels2 = pixels[len(pixels)//2:]

            pixels = (pixels1 & 0xff0000) | (pixels2 & 0xff)

        pygame.surfarray.blit_array(self.screen, pixels.reshape(self.size))
        if self.doom_mode:
            self.screen.blit(self.doom_hud, self.doom_rect)
        pygame.display.flip()

        if self.movie:
            self.screenshot(self.movie_dir, self.movie_index)
            self.movie_index += 1

    def loadlayer(self, layer):
        try:
            layergeometry = self.bvh_layers[layer]
        except KeyError:
            layergeometry = build(bvh_mesh(self.geometry, layer),8)
            self.bvh_layers[layer] = layergeometry

        self.gpu.load_geometry(layergeometry)
        self.update()

    def process_event(self, event):
        if event.type == MOUSEBUTTONDOWN:
            if event.button == 4:
                v = self.scale*np.cross(self.axis1,self.axis2)/10.0
                self.translate(v)

            elif event.button == 5:
                v = -self.scale*np.cross(self.axis1,self.axis2)/10.0
                self.translate(v)

            elif event.button == 1:
                mouse_position = pygame.mouse.get_rel()
                self.clicked = True

        elif event.type == MOUSEBUTTONUP:
            if event.button == 1:
                self.clicked = False

        elif event.type == MOUSEMOTION and self.clicked:
            movement = np.array(pygame.mouse.get_rel())

            if (movement == 0).all():
                return

            length = np.linalg.norm(movement)

            mouse_direction = movement[0]*self.axis2 - movement[1]*self.axis1
            mouse_direction /= np.linalg.norm(mouse_direction)

            if pygame.key.get_mods() & (KMOD_LSHIFT | KMOD_RSHIFT):
                v = -mouse_direction*self.scale*length/float(self.width)
                self.translate(v)
            else:
                phi = np.float32(2*np.pi*length/float(self.width))
                n = rotate(mouse_direction, np.pi/2, np.cross(self.axis1,self.axis2))

                if pygame.key.get_mods() & KMOD_LCTRL:
                    self.rotate_around_point(phi, n, self.point)
                else:
                    self.rotate(phi, n)

        elif event.type == KEYDOWN:
            if event.key == K_a:
                v = self.scale*self.axis2/10.0
                self.translate(v)

            elif event.key == K_d:
                v = -self.scale*self.axis2/10.0
                self.translate(v)

            elif event.key == K_w:
                v = self.scale*np.cross(self.axis1,self.axis2)/10.0
                self.translate(v)

            elif event.key == K_s:
                v = -self.scale*np.cross(self.axis1,self.axis2)/10.0
                self.translate(v)

            elif event.key == K_F6:
                self.clear_xyz_lookup()
                self.clear_image()
                self.source_position = self.point

            elif event.key == K_p:
                for i in range(100):
                    self.update_xyz_lookup(self.point)
                self.source_position = self.point

            elif event.key == K_F11:
                pygame.display.toggle_fullscreen()

            elif event.key == K_ESCAPE:
                self.done = True
                return

            elif event.key == K_PAGEDOWN:
                if self.currentlayer is None:
                    self.currentlayer = self.unique_bvh_layers[-1]
                else:
                    if self.currentlayer > np.min(self.unique_bvh_layers):
                        self.currentlayer -= 1
                    else:
                        self.currentlayer = None

                self.loadlayer(self.currentlayer)

            elif event.key == K_PAGEUP:
                if self.currentlayer is None:
                    self.currentlayer = self.unique_bvh_layers[0]
                else:
                    if self.currentlayer < np.max(self.unique_bvh_layers):
                        self.currentlayer += 1
                    else:
                        self.currentlayer = None

                self.loadlayer(self.currentlayer)                    

            elif event.key == K_F12:
                self.screenshot()

            elif event.key == K_F5:
                if not hasattr(self, 'rng_states_gpu'):
                    self.initialize_render()

                self.render = not self.render
                self.clear_image()
                self.update()

            elif event.key == K_F8:
                self.alpha = not self.alpha
                self.update()

            elif event.key == K_m:
                if self.movie:
                    encode_movie(self.movie_dir)
                    self.movie_dir = None
                    self.movie = False
                else:
                    self.movie_index = 0
                    self.movie_dir = tempfile.mkdtemp()
                    self.movie = True

        elif event.type == pygame.SYSWMEVENT and self.spnav:
            # Space Navigator controls
            spnav_event = self.spnav_module.spnav_x11_event(event.event)
            if spnav_event is None:
                return

            if spnav_event.type == self.spnav_module.SPNAV_EVENT_MOTION:
                if pygame.key.get_mods() & (KMOD_LSHIFT | KMOD_RSHIFT):
                    accelerate_factor = 2.0
                else:
                    accelerate_factor = 1.0
                #print 'raw:', spnav_event
                v1 = self.axis2
                v2 = self.axis1
                v3 = -np.cross(self.axis1,self.axis2)

                v = -v1*spnav_event.motion.x + v2*spnav_event.motion.y \
                    + v3*spnav_event.motion.z
                v *= self.scale / 5000.0 * accelerate_factor

                #print 'translate:', v
                self.translate(v, redraw=False)                    

                axis = v1*spnav_event.motion.rx + -v2*spnav_event.motion.ry \
                       + -v3*spnav_event.motion.rz

                # Zero all but non-max values
                #axis[~(np.abs(axis) == np.abs(axis).max())] = 0
                #axis[np.abs(axis) < 15] = 0
                if (axis != 0).any():
                    axis = axis.astype(float)
                    length = np.linalg.norm(axis)
                    angle = length * 0.0001 * accelerate_factor
                    axis /= length
                    #print 'rotate:', angle, axis
                    self.rotate_around_point(angle, axis, self.point, redraw=False)

                if self.render:
                    self.clear_image()

                self.update()
                pygame.event.clear(pygame.SYSWMEVENT)

            elif spnav_event.type == self.spnav_module.SPNAV_EVENT_BUTTON:
                if spnav_event.button.bnum == 0 and spnav_event.button.press:
                    if not hasattr(self, 'rng_states_gpu'):
                        self.initialize_render()

                    self.render = not self.render
                    self.clear_image()
                    self.update()
                    pygame.event.clear(pygame.SYSWMEVENT)

        elif event.type == pygame.QUIT:
            self.done = True
            return

    def run(self):
        self.init_gpu()

        if self.spnav:
            try:
                wm_info = pygame.display.get_wm_info()
                self.spnav_module.spnav_x11_open(wm_info['display'],
                                                 wm_info['window'])
                pygame.event.set_allowed(pygame.SYSWMEVENT)
                print 'Space Navigator support enabled.'
            except:
                self.spnav = False

        self.update()

        if self.enable3d:
            self.update()
        
        self.done = False
        self.clicked = False

        while not self.done:
            self.clock.tick(20)

            if self.render and not self.clicked and not pygame.event.peek(KEYDOWN):
                self.update()
            
            # Grab only last SYSWMEVENT (SpaceNav!) to avoid lagged controls
            for event in pygame.event.get(pygame.SYSWMEVENT)[-1:] + pygame.event.get():                
                self.process_event(event)

        if self.movie:
            encode_movie(self.movie_dir)

        pygame.display.quit()
        if self.spnav:
            self.spnav_module.spnav_close()

        del self.gpu

class EventViewer(Camera):
    def __init__(self, geometry, filename, **kwargs):
        Camera.__init__(self, geometry, **kwargs)

        import ROOT

        self.f = ROOT.TFile(filename)
        self.T = self.f.Get('T')
        self.T.GetEntry(0)
        self.nsolids = geometry.solid_id.max() + 1

    def color_hit_pmts(self):
        self.gpu.reset_colors()

        hit = np.empty(self.nsolids, np.int32)
        t = np.empty(self.nsolids, np.float32)
        q = np.empty(self.nsolids, np.float32)

        # self.nsolids has a weird data type that PyROOT doesn't understand
        self.T.ev.get_channels(int(self.nsolids), hit, t, q)

        # PyROOT prints warnings when we try to pass a bool array directly
        # so we convert afterward
        hit = hit.astype(np.bool)

        # Important: Compute range only with HIT channels
        solid_colors = map_to_color(q, range=(q[hit].min(),q[hit].max()))
        self.gpu.color_solids(hit, solid_colors)
        self.update()

    def process_event(self, event):
        if event.type == KEYDOWN:
            if event.key == K_PAGEUP:
                entry = self.T.GetReadEntry()
                if entry < self.T.GetEntries() - 1:
                    self.T.GetEntry(entry+1)
                    self.color_hit_pmts()
                    return

            elif event.key == K_PAGEDOWN:
                entry = self.T.GetReadEntry()
                if entry > 0:
                    self.T.GetEntry(entry-1)
                    self.color_hit_pmts()
                    return

        Camera.process_event(self, event)

if __name__ == '__main__':
    import optparse
    import inspect

    import chroma.solids
    import chroma.detectors
    import chroma.scenes
    from chroma.stl import mesh_from_stl

    parser = optparse.OptionParser('%prog filename.stl')
    parser.add_option('-b', '--bits', type='int', dest='bits',
                      help='bits for z-ordering space axes', default=10)
    parser.add_option('-r', '--resolution', dest='resolution',
                      help='specify resolution', default='1024,576')
    parser.add_option('--3d', action='store_true', dest='enable3d',
                      help='enable 3d', default=False)
    parser.add_option('-i', dest='io_file', default=None)
    options, args = parser.parse_args()

    if len(args) < 1:
        sys.exit(parser.format_help())

    size = [int(s) for s in options.resolution.split(',')]

    if os.path.exists(args[0]):
        root, ext = os.path.splitext(os.path.split(args[0])[1])

        if ext.lower() in ('.stl', '.bz2'):
            obj = mesh_from_stl(args[0])

    else:
        members = dict(inspect.getmembers(chroma.detectors) + inspect.getmembers(chroma.solids) + inspect.getmembers(chroma.scenes))

        buildable_lookup = {}
        for member in members.values():
            if inspect.isfunction(member) and \
                    hasattr(member, 'buildable') and member.buildable == True:
                buildable_lookup[member.identifier] = member

        if args[0] in buildable_lookup:
            obj = buildable_lookup[args[0]]
        else:
            raise Exception("can't find object %s" % args[0])

    geometry = build(obj, options.bits)
    if options.io_file is not None:
        camera = EventViewer(geometry, options.io_file, size=size, enable3d=options.enable3d)
    else:
        camera = Camera(geometry, size, enable3d=options.enable3d)

    camera.start()
    camera.join()