import numpy as np
from copy import copy
import time
import pycuda.driver as cuda
from pycuda.characterize import sizeof
from pycuda.compiler import SourceModule
from pycuda import gpuarray
import threading
import Queue
import src

class GPUThread(threading.Thread):
    def __init__(self, device_id, geometry, jobs, output, nblocks=64, max_nthreads=100000):
        threading.Thread.__init__(self)

        self.device_id = device_id
        self.geometry = copy(geometry)
        self.jobs = jobs
        self.output = output
        self.nblocks = nblocks
        self.max_nthreads = max_nthreads
        self._stop = threading.Event()

    def stop(self):
        self._stop.set()

    def stopped(self):
        return self._stop.is_set()

    def run(self):
        device = cuda.Device(self.device_id)
        context = device.make_context()
        module = SourceModule(src.kernel, options=['-I' + src.dir], no_extern_c=True, cache_dir=False)

        propagate = module.get_function('propagate')
        fill_uniform = module.get_function('fill_uniform')
        fill_uniform_sphere = module.get_function('fill_uniform_sphere')

        init_rng = module.get_function('init_rng')
        rng_states_gpu = cuda.mem_alloc(self.max_nthreads*sizeof('curandStateXORWOW', '#include <curand_kernel.h>'))
        init_rng(np.int32(self.max_nthreads), rng_states_gpu, np.int32(self.device_id), np.int32(0), block=(self.nblocks,1,1), grid=(self.max_nthreads//self.nblocks+1,1))

        self.geometry.load(module)

        daq_module = SourceModule(src.daq, options=[<
#!/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

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.05):
    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, spnav=False):
        Thread.__init__(self)
        self.geometry = geometry
        self.device_id = device_id
        self.size = size

        self.spnav = spnav

        if spnav:
            import spnav as spnav_module
            self.spnav_module = spnav_module

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

        self.kernels = CUDAFuncs(self.gpu.module, ['ray_trace', '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.75, (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.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.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:
            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_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:
            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)

        self.update()
        
        self.done = False
        self.clicked = False
        #shift = False
        #ctrl = 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)

    #@timeit
    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, (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)
    parser.add_option('--spnav', action='store_true', dest='spnav',
                      help='activate Space Navigator support', default=False)
    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, spnav=options.spnav)
    else:
        camera = Camera(geometry, size, spnav=options.spnav)
    camera.start()