#!/usr/bin/env python
import numpy as np

import pygame
from pygame.locals import *

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

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

def view(geometry, name=''):
    """
    Render `geometry` in a pygame window.

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

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

    pygame.init()
    size = width, height = 800, 600
    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_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]

    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)

    nblocks = 64

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

    def render():
        """Render the mesh and display to screen."""
        cuda_raytrace(np.int32(pixels.size), origins_gpu, directions_gpu, np.int32(geometry.node_map.size-1), np.int32(geometry.first_node), pixels_gpu, texrefs=texrefs, **gpu_kwargs)
        cuda.Context.synchronize()

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

    render()

    done = False
    clicked = False
    shift = False

    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.type == KEYUP:
                if event.key == K_LSHIFT or event.key == K_RSHIFT:
                    shift = False

if __name__ == '__main__':
    import os
    import sys
    import optparse

    from stl import *
    from materials import *
    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)
    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':
        geometry = Geometry()
        geometry.add_solid(Solid(read_stl(args[0]), vacuum, vacuum))
        geometry.build(options.bits)
        view(geometry, tail)
    else:
        import inspect

        members = dict(inspect.getmembers(detectors))

        if args[0] in members:
            if issubclass(members[args[0]], Geometry):
                geometry = members[args[0]]()
                geometry.build(options.bits)
                view(geometry, args[0])
            else:
                sys.exit("%s doesn't inherit from a Geometry object" % args[0])
        else:
            sys.exit("couldn't find a geometry with name %s" % args[0])