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

from gpu import GPU, CUDAFuncs
from tools import timeit
from transform import rotate

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 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

        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()

        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))
        pygame.display.flip()

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

    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_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)

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

        solid_ids = np.empty(len(self.T.ev.channel), np.uint32)
        t = np.empty(len(self.T.ev.channel), np.float32)
        q = np.empty(len(self.T.ev.channel), np.float32)

        for i, channel in enumerate(self.T.ev.channel):
            solid_ids[i] = channel.channel_id
            t[i] = channel.t
            q[i] = channel.q

        self.gpu.color_solids(solid_ids, map_to_color(t, range=(t.min(),t.mean())))
        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 solids
    import detectors
    import scenes
    from stl import mesh_from_stl
    from view import build

    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(detectors) + inspect.getmembers(solids) + inspect.getmembers(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()