import numpy as np
from itertools import product, count
from threading import Thread, Lock
import time
import datetime
import os
import sys
from color import map_wavelength

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 timeit(func):
    def f(*args, **kwargs):
        t0 = time.time()
        func(*args, **kwargs)
        elapsed = time.time() - t0
        print '%s elapsed in %s().' % (datetime.timedelta(seconds=elapsed), func.__name__)
    return f

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), film_size = (0.035, 0.024), focal_length=0.05):
    """
    Generate ray positions and directions from a pinhole camera facing the negative y direction.

    Args:
        - position: tuple,
            Position of the camera.
        - size: tuple, *optional*
            Pixel array shape.
        - film_size: tuple, *optional*
            Physical size of photographic film. Defaults to 35mm film size.
        - focal_length: float, *optional*
            Focal length of camera.
    """

    x = np.linspace(-film_size[0]/2, film_size[0]/2, size[0])
    z = np.linspace(-film_size[1]/2, film_size[1]/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, module, context, lock, size=(800,600)):
        Thread.__init__(self)
        self.geometry = geometry
        self.module = module
        self.context = context
        self.lock = lock

        self.size = size
        self.width, self.height = size

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

        with self.lock:
            self.context.push()
            self.ray_trace_kernel = self.module.get_function('ray_trace')
            self.rotate_kernel = self.module.get_function('rotate')
            self.rotate_around_point_kernel = self.module.get_function('rotate_around_point')
            self.translate_kernel = self.module.get_function('translate')
            self.update_xyz_lookup_kernel = self.module.get_function('update_xyz_lookup')
            self.update_xyz_image_kernel = self.module.get_function('update_xyz_image')
            self.process_image_kernel = self.module.get_function('process_image')
            self.init_rng_kernel = self.module.get_function('init_rng')
            self.context.pop()

        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)

        with self.lock:
            self.context.push()
            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.context.pop()

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

    @timeit
    def initialize_render(self):
        with self.lock:
            self.context.push()
            self.rng_states_gpu = cuda.mem_alloc(self.width*self.height*sizeof('curandStateXORWOW', '#include <curand_kernel.h>'))
            self.init_rng_kernel(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.context.pop()

        self.source_position = self.point

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

    def clear_xyz_lookup(self):
        with self.lock:
            self.context.push()
            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.context.pop()

        self.nlookup_calls = 0

    def update_xyz_lookup(self, source_position):
        with self.lock:
            self.context.push()
            for i in range(self.xyz_lookup1_gpu.size//(self.width*self.height)+1):
                self.update_xyz_lookup_kernel(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.update_xyz_lookup_kernel(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.update_xyz_lookup_kernel(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.context.pop()

        self.nlookup_calls += 1

    def clear_image(self):
        with self.lock:
            self.context.push()
            self.image_gpu.fill(gpuarray.vec.make_float3(0.0,0.0,0.0))
            self.context.pop()

        self.nimages = 0

    def update_image(self):
        with self.lock:
            self.context.push()
            self.update_xyz_image_kernel(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.update_xyz_image_kernel(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.update_xyz_image_kernel(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.context.pop()

        self.nimages += 1

    def process_image(self):
        with self.lock:
            self.context.push()
            self.process_image_kernel(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))
            self.context.pop()

    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):
        with self.lock:
            self.context.push()
            self.rotate_kernel(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.rotate_kernel(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)
            self.context.pop()

        if self.render:
            self.clear_image()

        self.update()

    def rotate_around_point(self, phi, n, point):
        with self.lock:
            self.context.push()
            self.rotate_around_point_kernel(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.rotate_kernel(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.context.pop()

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

        if self.render:
            self.clear_image()

        self.update()

    def translate(self, v):
        with self.lock:
            self.context.push()
            self.translate_kernel(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
            self.context.pop()

        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:
            with self.lock:
                self.context.push()
                self.ray_trace_kernel(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))
                self.context.pop()

        with self.lock:
            self.context.push()
            pygame.surfarray.blit_array(self.screen, self.pixels_gpu.get().reshape(self.size))
            pygame.display.flip()
            self.context.pop()

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

    def run(self):
        self.update()
        
        done = False
        clicked = False
        shift = False
        ctrl = False

        #current_layer = 0

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

            if self.render and not clicked and not pygame.event.peek(KEYDOWN):
                self.update()

            for event in pygame.event.get():
                if event.type == MOUSEBUTTONDOWN:
                    if event.button == 4:
                        v = self.scale*np.cross(self.axis1,self.axis2)/10.0
                        self.translate(v)

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

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

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

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

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

                    length = np.linalg.norm(movement)

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

                    if shift:
                        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 ctrl:
                            self.rotate_around_point(phi, n, self.point)
                        else:
                            self.rotate(phi, n)

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

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

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

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

                    if event.key == K_SPACE:
                        v = self.scale*self.axis2/10.0
                        self.translate(v)

                    # if event.key == K_LCTRL:
                    #     v = -self.scale*self.axis2/10.0
                    #     self.translate(v)

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

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

                    if event.key == K_LSHIFT or event.key == K_RSHIFT:
                        shift = True

                    if event.key == K_LCTRL or event.key == K_RCTRL:
                        ctrl = True

                    if event.key == K_ESCAPE:
                        done = True
                        break

                    #if event.key == K_PAGEUP and load_bvh:
                    #    try:
                    #        if current_layer+1 >= len(bvhg):
                    #            raise IndexError

                    #        geometry = bvhg[current_layer+1]
                    #        current_layer += 1

                    #        geometry.load(module, color=True)
                    #        update()
                    #    except IndexError:
                    #        print 'no further layers to view'

                    #if event.key == K_PAGEDOWN and load_bvh:
                    #    try:
                    #        if current_layer-1 < 0:
                    #            raise IndexError

                    #        geometry = bvhg[current_layer-1]
                    #        current_layer -= 1

                    #        geometry.load(module, color=True)
                    #        update()
                    #    except IndexError:
                    #        print 'no further layers to view'

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

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

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

                    if 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

                if event.type == KEYUP:
                    if event.key == K_LSHIFT or event.key == K_RSHIFT:
                        shift = False

                    if event.key == K_LCTRL or event.key == K_RCTRL:
                        ctrl = False

                if event.type == pygame.QUIT:
                    done = True
                    break

        if self.movie:
            encode_movie(self.movie_dir)

        pygame.display.quit()

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

    import solids
    import detectors
    import scenes
    from stl import mesh_from_stl
    import src
    from view import build

    from pycuda.compiler import SourceModule

    parser = optparse.OptionParser('%prog filename.stl')
    parser.add_option('-b', '--bits', type='int', dest='bits',
                      help='bits for z-ordering space axes', default=8)
    #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')
    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])

    from pycuda.tools import make_default_context

    cuda.init()
    context = make_default_context()
    print 'device %s' % context.get_device().name()

    module = SourceModule(src.kernel, options=['-I' + src.dir], no_extern_c=True, cache_dir=False)
    geometry = build(obj, options.bits)
    geometry.load(module)

    lock = Lock()

    camera = Camera(geometry, module, context, lock, size)

    context.pop()
    camera.start()