#!/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
from chroma.tools import timeit
from chroma.transform import rotate
from chroma.optics import vacuum, lambertian_surface

import pygame
from pygame.locals import *

from pycuda import gpuarray
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):
        Thread.__init__(self)
        self.geometry = geometry
        self.device_id = device_id
        self.size = size

        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.kernels = CUDAFuncs(self.gpu.module, ['ray_trace', 'ray_trace_alpha', 'rotate', 'rotate_around_point', 'translate', 'update_xyz_lookup', 'update_xyz_image', 'process_image', 'init_rng'])

        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([1,0,0], float)
        self.axis2 = np.array([0,0,1], float)

        origins, directions = get_rays(self.point, self.size)

        self.origins_gpu = gpuarray.to_gpu(origins.astype(np.float32).view(gpuarray.vec.float3))
        self.directions_gpu = gpuarray.to_gpu(directions.astype(np.float32).view(gpuarray.vec.float3))
        self.pixels_gpu = gpuarray.zeros(self.width*self.height, 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.width*self.height*sizeof('curandStateXORWOW', '#include <curand_kernel.h>'))
        self.kernels.init_rng(np.int32(self.width*self.height), self.rng_states_gpu, np.int32(0), np.int32(0), block=(self.nblocks,1,1), grid=(self.width*self.height//self.nblocks+1,1))
        self.xyz_lookup1_gpu = gpuarray.zeros(len(self.geometry.mesh.triangles), dtype=gpuarray.vec.float3)
        self.xyz_lookup2_gpu = gpuarray.zeros(len(self.geometry.mesh.triangles), dtype=gpuarray.vec.float3)
        self.image_gpu = gpuarray.zeros(self.width*self.height, dtype=gpuarray.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(gpuarray.vec.make_float3(0.0,0.0,0.0))
        self.xyz_lookup2_gpu.fill(gpuarray.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.width*self.height)+1):
            self.kernels.update_xyz_lookup(np.int32(self.width*self.height), np.int32(self.xyz_lookup1_gpu.size), np.int32(i*self.width*self.height), gpuarray.vec.make_float3(*source_position), self.rng_states_gpu, np.float32(685.0), gpuarray.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.width*self.height//self.nblocks+1,1))

        for i in range(self.xyz_lookup1_gpu.size//(self.width*self.height)+1):
            self.kernels.update_xyz_lookup(np.int32(self.width*self.height), np.int32(self.xyz_lookup1_gpu.size), np.int32(i*self.width*self.height), gpuarray.vec.make_float3(*source_position), self.rng_states_gpu, np.float32(545.0), gpuarray.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.width*self.height//self.nblocks+1,1))

        for i in range(self.xyz_lookup1_gpu.size//(self.width*self.height)+1):
            self.kernels.update_xyz_lookup(np.int32(self.width*self.height), np.int32(self.xyz_lookup1_gpu.size), np.int32(i*self.width*self.height), gpuarray.vec.make_float3(*source_position), self.rng_states_gpu, np.float32(445.0), gpuarray.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.width*self.height//self.nblocks+1,1))

        self.nlookup_calls += 1

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

        self.nimages = 0

    def update_image(self):
        self.kernels.update_xyz_image(np.int32(self.width*self.height), self.rng_states_gpu, self.origins_gpu, self.directions_gpu, np.float32(685.0), gpuarray.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.width*self.height//self.nblocks+1,1))

        self.kernels.update_xyz_image(np.int32(self.width*self.height), self.rng_states_gpu, self.origins_gpu, self.directions_gpu, np.float32(545.0), gpuarray.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.width*self.height//self.nblocks+1,1))

        self.kernels.update_xyz_image(np.int32(self.width*self.height), self.rng_states_gpu, self.origins_gpu, self.directions_gpu, np.float32(445.0), gpuarray.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.width*self.height//self.nblocks+1,1))

        self.nimages += 1

    def process_image(self):
        self.kernels.process_image(np.int32(self.width*self.height), self.image_gpu, self.pixels_gpu, np.int32(self.nimages), block=(self.nblocks,1,1), grid=((self.width*self.height)//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.kernels.rotate(np.int32(self.pixels_gpu.size), self.origins_gpu, np.float32(phi), gpuarray.vec.make_float3(*n), block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks+1,1))
        self.kernels.rotate(np.int32(self.pixels_gpu.size), self.directions_gpu, np.float32(phi), gpuarray.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.render:
            self.clear_image()

        self.update()

    def rotate_around_point(self, phi, n, point, redraw=True):
        self.kernels.rotate_around_point(np.int32(self.origins_gpu.size), self.origins_gpu, np.float32(phi), gpuarray.vec.make_float3(*n), gpuarray.vec.make_float3(*point), block=(self.nblocks,1,1), grid=(self.origins_gpu.size//self.nblocks+1,1))
        self.kernels.rotate(np.int32(self.directions_gpu.size), self.directions_gpu, np.float32(phi), gpuarray.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 redraw:
            if self.render:
                self.clear_image()

            self.update()

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

        self.point += v

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

            self.update()

    def update(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.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.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))

        pygame.surfarray.blit_array(self.screen, self.pixels_gpu.get().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.axis1 + movement[1]*self.axis2
            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.axis1/10.0
                self.translate(v)

            elif event.key == K_d:
                v = -self.scale*self.axis1/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_SPACE:
                v = self.scale*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.axis1
                v2 = self.axis2
                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()
        
        self.done = False
        self.clicked = False

        #current_layer = 0

        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

            if 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('-l', action='store_true', dest='load_bvh',
    #                  help='load bounding volumes', default=False)
    parser.add_option('-r', '--resolution', dest='resolution',
                      help='specify resolution', default='800,600')
    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)
    else:
        camera = Camera(geometry, size)
    camera.start()