#!/usr/bin/env python
import numpy as np
from stl import *
from geometry import *
from materials import *
from camera import *
from gpu import *
from transform import *

import pygame
from pygame.locals import *

def view(geometry, name=''):
    mesh = geometry.mesh
    lower_bound = np.array([np.min(mesh[:,:,0]), np.min(mesh[:,:,1]), np.min(mesh[:,:,2])])
    upper_bound = np.array([np.max(mesh[:,:,0]), np.max(mesh[:,:,1]), np.max(mesh[:,:,2])])

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

    gpu = GPU()
    gpu.load_geometry(geometry)

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

    camera = Camera(size)
    camera.position((0, upper_bound[1]*5, np.mean([lower_bound[2], upper_bound[2]])))

    origin, direction = camera.get_rays()

    pixels = np.empty(width*height, dtype=np.int32)
    states = np.empty(width*height, dtype=np.int32)

    origin_gpu = cuda.to_device(make_vector(origin))
    direction_gpu = cuda.to_device(make_vector(direction))

    pixels_gpu = cuda.to_device(pixels)
    states_gpu = cuda.to_device(states)

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

    def render():
        gpu.call(np.int32(pixels.size), origin_gpu, direction_gpu, np.int32(geometry.first_leaf), states_gpu, pixels_gpu, **gpu_kwargs)

        cuda.memcpy_dtoh(pixels, pixels_gpu)

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

    render()

    done = False
    clicked = False
    to_origin = np.array([0,-1,0])
    while not done:
        for event in pygame.event.get():
            if event.type == MOUSEBUTTONDOWN:
                if event.button == 4:
                    cuda_translate(np.int32(pixels.size), origin_gpu, gpuarray.vec.make_float3(*(scale*to_origin)), **gpu_kwargs)
                    render()

                if event.button == 5:
                    cuda_translate(np.int32(pixels.size), origin_gpu, gpuarray.vec.make_float3(*(-scale*to_origin)), **gpu_kwargs)
                    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 = pygame.mouse.get_rel()[0]/float(width)

                if movement == 0:
                    continue

                cuda_rotate(np.int32(pixels.size), origin_gpu, np.float32(movement*2*np.pi), gpuarray.vec.make_float3(0,0,1), **gpu_kwargs)

                cuda_rotate(np.int32(pixels.size), direction_gpu, np.float32(movement*2*np.pi), gpuarray.vec.make_float3(0,0,1), **gpu_kwargs)

                to_origin = rotate(to_origin, movement*2*np.pi, (0,0,1))

                render()

            if event.type == KEYUP or event.type == KEYDOWN:
                if event.key == K_ESCAPE:
                    done = True
                    break

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

    parser = optparse.OptionParser('%prog filename.stl')
    parser.add_option('-b', '--bits', type='int', dest='bits', help='number of bits for z ordering space axes', default=4)
    options, args = parser.parse_args()

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

    geometry = Geometry()
    geometry.add_solid(Solid(read_stl(args[0]), vacuum, vacuum))
    geometry.build(options.bits)

    view(geometry, os.path.split(args[0])[-1])