#!/usr/bin/env python
import os
import sys
import numpy as np
import pickle
import tempfile

import pygame
from pygame.locals import *

import src
from camera import *
from geometry import *
from transform import *
from optics import *

import time

from pycuda import autoinit
from pycuda.compiler import SourceModule
from pycuda import gpuarray
from pycuda.characterize import sizeof

def to_geometry(viewable, bits):
    if isinstance(viewable, Geometry):
        geometry = viewable
        geometry.build(bits)
    elif isinstance(viewable, Solid):
        geometry = Geometry()
        geometry.add_solid(viewable)
        geometry.build(bits)
    elif isinstance(viewable, Mesh):
        geometry = Geometry()
        geometry.add_solid(Solid(viewable, vacuum, vacuum, surface=lambertian_surface))
        geometry.build(bits)
    else:
        raise Exception("can't convert %s to a geometry" % viewable)

    return geometry

def box(lower_bound, upper_bound):
    dx, dy, dz = upper_bound - lower_bound

    vertices = np.array([lower_bound,
                         lower_bound + (dx,0,0),
                         lower_bound + (dx,dy,0),
                         lower_bound + (0,dy,0),
                         lower_bound + (0,0,dz),
                         lower_bound + (dx,0,dz),
                         lower_bound + (dx,dy,dz),
                         lower_bound + (0,dy,dz)])

    triangles = np.empty((12,3), dtype=np.int32)
    # bottom
    triangles[0] = (1,3,2)
    triangles[1] = (1,4,3)
    # top
    triangles[2] = (5,7,6)
    triangles[3] = (5,8,7)
    # left
    triangles[4] = (5,1,2)
    triangles[5] = (5,2,6)
    # right
    triangles[6] = (3,4,8)
    triangles[7] = (3,8,7)
    # front
    triangles[8] = (2,3,7)
    triangles[9] = (2,7,6)
    # back
    triangles[10] = (1,5,8)
    triangles[11] = (1,8,4)

    triangles -= 1

    return Mesh(vertices, triangles)

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)

    mesh = box(lower_bounds[0], upper_bounds[0])

    for lower_bound, upper_bound in zip(lower_bounds, upper_bounds)[1:]:
        mesh += box(lower_bound, upper_bound)

    return mesh

def view(viewable, name='', bits=8, load_bvh=False, camera_view=None, make_movie=False):
    """
    Render `viewable` in a pygame window.

    Movement:
        - zoom: scroll the mouse wheel
        - rotate: click and drag the mouse
        - move: shift+click and drag the mouse
    """

    geometry = to_geometry(viewable, bits)

    if load_bvh:
        print 'loading bvh...'
        bvhg = []

        bvhg.append(geometry)

        for layer in sorted(np.unique(geometry.layers)):
            print 'building layer %i' % layer
            bvhg.append(to_geometry(bvh_mesh(geometry, layer), bits))

    lower_bound = np.array([np.min(geometry.mesh[:][:,:,0]),
                            np.min(geometry.mesh[:][:,:,1]),
                            np.min(geometry.mesh[:][:,:,2])])

    upper_bound = np.array([np.max(geometry.mesh[:][:,:,0]),
                            np.max(geometry.mesh[:][:,:,1]),
                            np.max(geometry.mesh[:][:,:,2])])

    scale = np.linalg.norm(upper_bound-lower_bound)

    print 'device %s' % autoinit.device.name()

    module = SourceModule(src.kernel, options=['-I' + src.dir],
                          no_extern_c=True, cache_dir=False)
    #geometry.load(module, color=True)
    geometry.load(module, color=False)
    #cuda_raytrace = module.get_function('ray_trace')
    cuda_build_rgb_lookup = module.get_function('build_rgb_lookup')
    cuda_render = module.get_function('render')
    cuda_rotate = module.get_function('rotate')
    cuda_translate = module.get_function('translate')

    pygame.init()
    size = width, height = 800, 600

    init_rng = module.get_function('init_rng')
    rng_states_gpu = cuda.mem_alloc(len(geometry.mesh.triangles)*sizeof('curandStateXORWOW', '#include <curand_kernel.h>'))
    init_rng(np.int32(len(geometry.mesh.triangles)), rng_states_gpu, np.int32(0), np.int32(0), block=(64,1,1), grid=(len(geometry.mesh.triangles)//64+1,1))

    diagonal = np.linalg.norm(upper_bound-lower_bound)
    #source_position = [0, diagonal*1.75, (lower_bound[2]+upper_bound[2])/2]
    source_position = [0, 0, upper_bound[2]+5.0]

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

    source_directions = np.mean(geometry.mesh[:], axis=1) - source_position
    source_directions_float3 = np.empty(source_positions_gpu.size, dtype=gpuarray.vec.float3)
    source_directions_float3['x'] = source_directions[:,0]
    source_directions_float3['y'] = source_directions[:,1]
    source_directions_float3['z'] = source_directions[:,2]
    source_directions_gpu = cuda.to_device(source_directions_float3)

    if (len(geometry.mesh.triangles) < source_positions_gpu.size):
        print width*height
        print source_positions_gpu.size
        raise Exception('not enough rng_states')

    rgb_lookup_gpu = gpuarray.zeros(source_positions_gpu.size, dtype=gpuarray.vec.float3)

    max_steps = 10
    rgb_runs = 1000

    cuda_build_rgb_lookup(np.int32(source_positions_gpu.size), rng_states_gpu, source_positions_gpu, source_directions_gpu, np.int32(geometry.node_map.size-1), np.int32(geometry.first_node), rgb_lookup_gpu, np.int32(rgb_runs), np.int32(max_steps), block=(64,1,1), grid=(source_positions_gpu.size//64+1,1))

    rgb_lookup = rgb_lookup_gpu.get()
    rgb_lookup['x'] /= rgb_runs
    rgb_lookup['y'] /= rgb_runs
    rgb_lookup['z'] /= rgb_runs
    rgb_lookup_gpu = cuda.to_device(rgb_lookup)


    print rgb_lookup












    if camera_view is None:
        camera = Camera(size)

        diagonal = np.linalg.norm(upper_bound-lower_bound)

        point = np.array([0, diagonal*1.75, (lower_bound[2]+upper_bound[2])/2])
        axis1 = np.array([1,0,0], dtype=np.double)
        axis2 = np.array([0,0,1], dtype=np.double)

        camera.position(point)

        origins, directions = camera.get_rays()

        origins_float3 = np.empty(origins.shape[0], dtype=gpuarray.vec.float3)
        origins_float3['x'] = origins[:,0]
        origins_float3['y'] = origins[:,1]
        origins_float3['z'] = origins[:,2]

        directions_float3 = np.empty(directions.shape[0], dtype=gpuarray.vec.float3)
        directions_float3['x'] = directions[:,0]
        directions_float3['y'] = directions[:,1]
        directions_float3['z'] = directions[:,2]
    else:
        f = open(camera_view, 'rb')
        origins_float3 = pickle.load(f)
        directions_float3 = pickle.load(f)
        point = pickle.load(f)
        axis1 = pickle.load(f)
        axis2 = pickle.load(f)
        f.close()

    origins_gpu = cuda.to_device(origins_float3)
    directions_gpu = cuda.to_device(directions_float3)

    pixels = np.empty(width*height, dtype=np.int32)
    pixels_gpu = cuda.to_device(pixels)

    nruns = 5



    screen = pygame.display.set_mode(size)
    pygame.display.set_caption(name)






    nblocks = 64

    gpu_kwargs = {'block': (nblocks,1,1), 'grid':(pixels.size/nblocks+1,1)}

    if make_movie:
        tempdir = tempfile.mkdtemp()

    def screenshot(dir=''):
        if name == '':
            root, ext = 'screenshot', 'png'
        else:
            root, ext = name, 'png'

        filename = os.path.join(dir,'.'.join([root, ext]))

        i = 1
        while os.path.exists(filename):
            filename = os.path.join(dir,'.'.join([root + str(i), ext]))
            i += 1

        pygame.image.save(screen, filename)
        print 'image saved to %s' % filename

    def render():
        """Render the mesh and display to screen."""
        t0 = time.time()
        cuda_render(np.int32(pixels.size), rng_states_gpu, origins_gpu, directions_gpu, np.int32(geometry.node_map.size-1), np.int32(geometry.first_node), rgb_lookup_gpu, np.int32(nruns), pixels_gpu, np.int32(max_steps), **gpu_kwargs)
        cuda.Context.synchronize()
        elapsed = time.time() - t0

        print 'elapsed %f sec' % elapsed

        cuda.memcpy_dtoh(pixels, pixels_gpu)
        pygame.surfarray.blit_array(screen, pixels.reshape(size))
        pygame.display.flip()

        if make_movie:
            screenshot(tempdir)

    render()

    done = False
    clicked = False
    shift = False

    current_layer = 0

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

                    cuda_translate(np.int32(pixels.size), origins_gpu, gpuarray.vec.make_float3(*v), **gpu_kwargs)

                    point += v

                    render()

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

                    cuda_translate(np.int32(pixels.size), origins_gpu, gpuarray.vec.make_float3(*v), **gpu_kwargs)

                    point += v

                    render()

                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]*axis1 + movement[1]*axis2
                mouse_direction /= np.linalg.norm(mouse_direction)

                if shift:
                    v = mouse_direction*scale*length/float(width)

                    cuda_translate(np.int32(pixels.size), origins_gpu, gpuarray.vec.make_float3(*v), **gpu_kwargs)

                    point += v

                    render()
                else:
                    phi = np.float32(2*np.pi*length/float(width))
                    n = rotate(mouse_direction, np.pi/2, \
                                   -np.cross(axis1,axis2))

                    cuda_rotate(np.int32(pixels.size), origins_gpu, phi, gpuarray.vec.make_float3(*n), **gpu_kwargs)

                    cuda_rotate(np.int32(pixels.size), directions_gpu, phi, gpuarray.vec.make_float3(*n), **gpu_kwargs)

                    point = rotate(point, phi, n)
                    axis1 = rotate(axis1, phi, n)
                    axis2 = rotate(axis2, phi, n)

                    render()

            if event.type == KEYDOWN:
                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:
                    try:
                        if current_layer+1 >= len(bvhg):
                            raise IndexError

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

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

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

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

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

                if event.key == K_F11:
                    if name == '':
                        root, ext = 'camera_view', 'pkl'
                    else:
                        root, ext = name + '_camera_view', 'pkl'

                    filename = '.'.join([root, ext])

                    i = 1
                    while os.path.exists(filename):
                        filename = '.'.join([root + str(i), ext])
                        i += 1

                    f = open(filename, 'wb')
                    cuda.memcpy_dtoh(origins_float3, origins_gpu)
                    cuda.memcpy_dtoh(directions_float3, directions_gpu)
                    pickle.dump(origins_float3, f)
                    pickle.dump(directions_float3, f)
                    pickle.dump(point, f)
                    pickle.dump(axis1, f)
                    pickle.dump(axis2, f)
                    f.close()

                    print 'camera view saved to %s' % filename

                if event.key == K_F12:
                    screenshot()

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

    pygame.display.quit()

    if make_movie:
        if name == '':
            root, ext = 'movie', 'tgz'
        else:
            root, ext = name + '_movie', 'tgz'

        filename = '.'.join([root, ext])

        i=1
        while os.path.exists(filename):
            filename = '.'.join([root + str(i), ext])
            i += 1

        from subprocess import call
        call(['tar', '-czvf', filename, tempdir])

        os.removedirs(tempdir)

if __name__ == '__main__':
    import optparse

    from stl import mesh_from_stl

    import solids
    import detectors

    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('-v', '--camera_view', dest='camera_view',
                      help='load a camera view', default=None)
    parser.add_option('-m', '--movie', action='store_true', dest='movie',
                      help='create a movie', default=False)
    options, args = parser.parse_args()

    if len(args) < 1:
        sys.exit(parser.format_help())

    head, tail = os.path.split(args[0])
    root, ext = os.path.splitext(tail)

    if ext.lower() == '.stl':
        view(mesh_from_stl(args[0]), root, options.bits, options.load_bvh, options.camera_view, options.movie)
    else:
        import inspect

        members = dict(inspect.getmembers(detectors) + inspect.getmembers(solids))

        if args[0] in members:
            view(members[args[0]], args[0], options.bits, options.load_bvh, options.camera_view, options.movie)
        else:
            sys.exit("couldn't find object %s" % args[0])