#!/usr/bin/env python
import os
import sys
import time
import numpy as np
import inspect

import pygame
from pygame.locals import *

import src
from camera import Camera
from geometry import Mesh, Solid, Geometry
from transform import rotate

from pycuda import autoinit
from pycuda.compiler import SourceModule
from pycuda import gpuarray
import pycuda.driver as cuda

def buildable(identifier):
    """
    Create a decorator which tags a function as buildable and assigns the
    identifying string `identifier`.

    Example:
        >>> @buildable('my_sphere')
        >>> def build_my_sphere():
        >>>     g = Geometry()
        >>>     g.add_solid(Solid(sphere(), vacuum, water))
        >>>     return g
    """
    def tag_as_buildable(func):
        func.buildable = True
        func.identifier = identifier
        return func
    return tag_as_buildable

def screenshot(screen, name='', dir='', index=0):
    """Take a screenshot of `screen`."""
    if name == '':
        root, ext = 'screenshot', 'png'
    else:
        root, ext = name, 'png'

    i = index

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

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

def build(obj, bits):
    """Construct and build a geometry from `obj`."""
    if inspect.isfunction(obj):
        try:
            if obj.buildable:
                obj = obj()
        except AttributeError:
            raise Exception('function %s is not buildable.' % obj.__name__)

    if isinstance(obj, Geometry):
        geometry = obj
    elif isinstance(obj, Solid):
        geometry = Geometry()
        geometry.add_solid(obj)
    elif isinstance(obj, Mesh):
        geometry = Geometry()
        geometry.add_solid(Solid(obj))
    else:
        raise Exception('cannot build type %s' % type(obj))

    geometry.build(bits)

    return geometry

def box(lower_bound, upper_bound):
    """
    Return a mesh of the box defined by the opposing corners `lower_bound`
    and `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, size=(800,600), name='', bits=8, load_bvh=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 = build(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(build(bvh_mesh(geometry, layer), bits))

    lower_bound, upper_bound = geometry.mesh.get_bounds()

    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)
    cuda_raytrace = module.get_function('ray_trace')
    cuda_rotate = module.get_function('rotate')
    cuda_translate = module.get_function('translate')

    pygame.init()
    width, height = size
    screen = pygame.display.set_mode(size)
    pygame.display.set_caption(name)

    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_gpu = gpuarray.to_gpu(origins.astype(np.float32).view(gpuarray.vec.float3))

    directions_gpu = gpuarray.to_gpu(directions.astype(np.float32).view(gpuarray.vec.float3))

    pixels_gpu = gpuarray.zeros(width*height, dtype=np.int32)

    nblocks = 64

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

    def update():
        """Render the mesh and display to screen."""
        t0 = time.time()
        cuda_raytrace(np.int32(pixels_gpu.size), origins_gpu, directions_gpu, np.int32(geometry.node_map.size-1), np.int32(geometry.first_node), pixels_gpu, block=(nblocks,1,1), grid=(pixels_gpu.size//nblocks+1,1))
        cuda.Context.synchronize()
        elapsed = time.time() - t0

        #print 'elapsed %f sec' % elapsed

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

    update()

    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_gpu.size), origins_gpu, gpuarray.vec.make_float3(*v), block=(nblocks,1,1), grid=(pixels_gpu.size//nblocks+1,1))

                    point += v

                    update()

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

                    cuda_translate(np.int32(pixels_gpu.size), origins_gpu, gpuarray.vec.make_float3(*v), block=(nblocks,1,1), grid=(pixels_gpu.size//nblocks+1,1))

                    point += v

                    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]*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_gpu.size), origins_gpu, gpuarray.vec.make_float3(*v), block=(nblocks,1,1), grid=(pixels_gpu.size//nblocks+1,1))

                    point += v

                    update()
                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_gpu.size), origins_gpu, phi, gpuarray.vec.make_float3(*n), block=(nblocks,1,1), grid=(pixels_gpu.size//nblocks+1,1))

                    cuda_rotate(np.int32(pixels_gpu.size), directions_gpu, phi, gpuarray.vec.make_float3(*n), block=(nblocks,1,1), grid=(pixels_gpu.size//nblocks+1,1))

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

                    update()

            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 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:
                    screenshot(screen, name)

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

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

    pygame.display.quit()

if __name__ == '__main__':
    import optparse

    from stl import mesh_from_stl

    import solids
    import detectors
    import scenes

    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]):
        head, tail = os.path.split(args[0])
        root, ext = os.path.splitext(tail)

        if ext.lower() == '.stl':
            view(mesh_from_stl(args[0]), size, root, options.bits, options.load_bvh)
    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:
            view(buildable_lookup[args[0]], size, args[0], options.bits, options.load_bvh)
        else:
            sys.exit("couldn't find object %s" % args[0])