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

from transform import rotate

import pygame
from pygame.locals import *

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

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.translate_kernel = self.module.get_function('translate')
            self.build_rgb_lookup_kernel = self.module.get_function('build_rgb_lookup')
            self.render_kernel = self.module.get_function('render')
            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)

        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.render = False

    def build_rgb_lookup(self, source_position=(0,0,0)):
        with self.lock:
            self.context.push()
            print '\ninitializing random number states.'
            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.context.synchronize()
            print 'done.'

            self.source_positions_gpu = gpuarray.empty(len(geometry.mesh.triangles), dtype=gpuarray.vec.float3)
            self.source_positions_gpu.fill(gpuarray.vec.make_float3(*source_position))

            source_directions = np.mean(self.geometry.mesh[:], axis=1) - source_position
            self.source_directions_gpu = gpuarray.to_gpu(source_directions.astype(np.float32).view(gpuarray.vec.float3))

            self.rgb_lookup1_gpu = gpuarray.zeros(self.source_positions_gpu.size, dtype=gpuarray.vec.float3)
            self.rgb_lookup2_gpu = gpuarray.zeros(self.source_positions_gpu.size, dtype=gpuarray.vec.float3)

            self.max_steps = 10
            rgb_runs = 100
            
            print 'building rgb lookup.'
            self.build_rgb_lookup_kernel(np.int32(self.source_positions_gpu.size), self.source_positions_gpu, self.source_directions_gpu, self.rgb_lookup1_gpu, self.rgb_lookup2_gpu, np.uint64(np.random.randint(np.iinfo(np.int64).max)), np.int32(rgb_runs), np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.source_positions_gpu.size//self.nblocks+1,1))
            self.context.synchronize()
            print 'done.'

            rgb_lookup1 = self.rgb_lookup1_gpu.get().view(np.float32)
            rgb_lookup1 /= rgb_runs
            rgb_lookup1[rgb_lookup1 > 1.0] = 1.0
            self.rgb_lookup1_gpu.set(rgb_lookup1.view(gpuarray.vec.float3))

            rgb_lookup2 = self.rgb_lookup2_gpu.get().view(np.float32)
            rgb_lookup2 /= rgb_runs
            rgb_lookup2[rgb_lookup2 > 1.0] = 1.0
            self.rgb_lookup2_gpu.set(rgb_lookup2.view(gpuarray.vec.float3))

            self.nimages = 0
            self.rgb_image = np.zeros((self.pixels_gpu.size, 3), float)
            self.r_gpu = gpuarray.zeros(self.pixels_gpu.size, dtype=np.float32)
            self.g_gpu = gpuarray.zeros(self.pixels_gpu.size, dtype=np.float32)
            self.b_gpu = gpuarray.zeros(self.pixels_gpu.size, dtype=np.float32)

            self.context.pop()

    def clear_image(self):
        try:
            self.rgb_image.fill(0.0)
            self.nimages = 0
        except AttributeError:
            pass

    def acquire_image(self):
        if not hasattr(self, 'rng_states_gpu'):
            self.build_rgb_lookup()

        with self.lock:
            self.context.push()
            self.render_kernel(np.int32(self.pixels_gpu.size), self.rng_states_gpu, self.origins_gpu, self.directions_gpu, self.r_gpu, self.g_gpu, self.b_gpu, self.rgb_lookup1_gpu, self.rgb_lookup2_gpu, np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks+1,1))
            self.rgb_image[:,0] += self.r_gpu.get()
            self.rgb_image[:,1] += self.g_gpu.get()
            self.rgb_image[:,2] += self.b_gpu.get()
            self.nimages += 1
            self.context.pop()

    def render_image(self):
        scaled_rgb_image = ((self.rgb_image/self.nimages)*255).astype(np.int32)

        image = scaled_rgb_image[:,0] << 16 | scaled_rgb_image[:,1] << 8 | scaled_rgb_image[:,2]

        pygame.surfarray.blit_array(self.screen, image.reshape(self.size))
        pygame.display.flip()

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

        for i in count():
            filename = os.path.join(dir, '.'.join([root + str(i).zfill(4), 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, 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, 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()

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

    def update(self):
        if self.render:
            self.acquire_image()
            self.render_image()
            return

        with self.lock:
            self.context.push()
            t0 = time.time()
            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.synchronize()
            elapsed = time.time() - t0

            print '\relapsed %f sec.' % elapsed,
            sys.stdout.flush()

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

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

        current_layer = 0

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

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

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

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

                    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)
                        self.clear_image()
                        self.update()
                    else:
                        phi = np.float32(2*np.pi*length/float(self.width))
                        n = rotate(mouse_direction, np.pi/2, -np.cross(self.axis1,self.axis2))

                        self.rotate(phi, n)
                        self.clear_image()
                        self.update()

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

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

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

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

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

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

                    if event.key == K_LSHIFT or event.key == K_RSHIFT:
                        shift = 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:
                        self.render = not self.render
                        self.update()

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

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

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