import numpy as np
from solid import Solid
from geometry import Geometry
from mesh import mesh_from_stl
from sample import uniform_sphere
from pycuda import autoinit
from pycuda.compiler import SourceModule
from pycuda import gpuarray
import pycuda.driver as cuda
import src
import models
from materials import *

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

lens_mesh = mesh_from_stl(models.dir + '/lens.stl')
lens_solid = Solid(0, lens_mesh, material1=pmt_glass, material2=pmt_vacuum)

film_mesh = mesh_from_stl(models.dir + '/film.stl')
film_solid = Solid(0, film_mesh, displacement=(-5,0,0), material1=h2o, material2=pmt_vacuum)

sphere_mesh = mesh_from_stl(models.dir + '/sphere.stl')
sphere_solid = Solid(0, sphere_mesh, displacement=(10,0,0), material1=h2o, material2=pmt_vacuum, surface=aluminum)

geometry = Geometry()
geometry.add_solid(lens_solid)
geometry.add_solid(film_solid)
geometry.add_solid(sphere_solid)
geometry.build(bits=8)

module = SourceModule(src.kernel, options=['-I' + src.dir], no_extern_c=True, cache_dir=False)
geometry.load(module)
cuda_propagate = module.get_function('propagate')

nphotons = 100

positions = np.tile((10,5,0), nphotons).reshape(nphotons,3)
positions_float3 = np.empty(positions.shape[0], dtype=gpuarray.vec.float3)
positions_float3['x'] = positions[:,0]
positions_float3['y'] = positions[:,1]
positions_float3['z'] = positions[:,2]
positions_gpu = cuda.to_device(positions_float3)

directions = uniform_sphere(nphotons)
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]
directions_gpu = cuda.to_device(directions_float3)

wavelengths = np.random.uniform(200, 800, nphotons).astype(np.float32)
wavelengths_gpu = cuda.to_device(wavelengths)

times = np.tile(0, nphotons).astype(np.float32)
times_gpu = cuda.to_device(times)

polarizations = uniform_sphere(nphotons)
polarizations_float3 = np.empty(polarizations.shape[0], dtype=gpuarray.vec.float3)
polarizations_float3['x'] = polarizations[:,0]
polarizations_float3['y'] = polarizations[:,1]
polarizations_float3['z'] = polarizations[:,2]
polarizations_gpu = cuda.to_device(polarizations_float3)

last_hit_triangles = np.empty(nphotons, dtype=np.int32)
last_hit_triangles_gpu = cuda.to_device(last_hit_triangles)

states = np.empty(nphotons, dtype=np.int32)
states_gpu = cuda.to_device(states)

nblocks = 64

gpu_kwargs = {'block': (nblocks,1,1), 'grid': (nphotons/nblocks+1,1)}

cuda_propagate(np.int32(nphotons), positions_gpu, directions_gpu, wavelengths_gpu, polarizations_gpu, times_gpu, states_gpu, last_hit_triangles_gpu, np.int32(geometry.node_map.size-1), np.int32(geometry.first_node), **gpu_kwargs)

cuda.memcpy_dtoh(states, states_gpu)

print states




#from view import view

#view(geometry)