#!/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 get_rays
from geometry import Mesh, Solid, Geometry
from transform import rotate
from optics import *
from gpu import *

#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, vacuum, vacuum, surface=lambertian_surface, color=0x99ffffff))
    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
    """

    gpu = GPU()

    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)

    #from pycuda import autoinit

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

    module = gpu.module

    #module = SourceModule(src.kernel, options=['-I' + src.dir], no_extern_c=True)#, cache_dir=False)
    #geometry.load(module)
    gpu.load_geometry(geometry)
    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 = get_rays(point, size)

    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

    def update():
        """Render the mesh and display to screen."""
        t0 = time.time()
        cuda_raytrace(np.int32(pixels_gpu.size), origins_gpu, directions_gpu, 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

                        gpu.load_geometry(geometry)
                        #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

                        gpu.load_geometry(geometry)
                        #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() in ('.stl', '.bz2'):
            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])