import sys
import time
import numpy as np
from solid import Solid
from geometry import Geometry
from mesh import mesh_from_stl
from sample import uniform_sphere, flashlight
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 *
from transform import *
from itertoolset import roundrobin
from view import view
import matplotlib.pyplot as plt

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

nphotons, nblocks = 1000000, 64

start_position = (5.0,0.0,0.0)

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

lens_mesh = mesh_from_stl(models.dir + '/lens.stl')
lens_mesh.vertices /= 1000
lens_solid = Solid(0, lens_mesh, material1=acrylic_sno, material2=vacuum, color=0x00ff0000)

box_mesh = mesh_from_stl(models.dir + '/box.stl')
box_mesh.vertices /= 1000
box_solid = Solid(0, box_mesh, material1=vacuum, material2=vacuum, surface=black_surface, color=0x000000ff)

film_mesh = mesh_from_stl(models.dir + '/film.stl')
film_mesh.vertices /= 10
film_solid = Solid(1, film_mesh, material1=vacuum, material2=vacuum, surface=black_surface, displacement=(-0.3,0,0), color=0x0000ff00)

pmt_mesh1 = mesh_from_stl(models.dir + '/hamamatsu_12inch.stl')
pmt_mesh1.vertices /= 1000
pmt_mesh2 = mesh_from_stl(models.dir + '/hamamatsu_12inch.stl')
pmt_mesh2.vertices /= 1000
pmt_mesh2.vertices *= 0.95
pmt_solid1 = Solid(3, pmt_mesh1, displacement=(2.0,0.0,0.0), material1=glass, material2=vacuum)
pmt_solid2 = Solid(4, pmt_mesh2, displacement=(2.0,0.0,0.0), material1=vacuum, material2=glass, surface=lambertian_surface)

sphere_mesh = mesh_from_stl(models.dir + '/sphere.stl')
sphere_mesh.vertices *= 100
sphere_solid = Solid(2, sphere_mesh, material1=vacuum, material2=vacuum, surface=black_surface)

geometry = Geometry()#[lens_solid1, lens_solid2, film_solid, sphere_solid])
geometry.add_solid(lens_solid)
geometry.add_solid(film_solid)
geometry.add_solid(box_solid)
#geometry.add_solid(pmt_solid1)
#geometry.add_solid(pmt_solid2)
#geometry.add_solid(sphere_solid)

#from detectors import LBNE
#geometry = LBNE()

texrefs = geometry.build(bits=8)

#view(geometry)

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

init_rng(np.int32(nphotons), np.int32(0), np.int32(0), **gpu_kwargs)

positions = np.tile(start_position, 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 = flashlight(np.pi/256, (-1,0,0), size=nphotons)
#directions = flashlight(np.pi/8, (0,-1,0), size=nphotons)
#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(400, 700, 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.ones(nphotons, dtype=np.int32)
last_hit_triangles_gpu = cuda.to_device(last_hit_triangles)

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

x = np.empty((nphotons, 10, 3))

for i in range(10):
    x[:,i,0] = positions_float3['x']
    x[:,i,1] = positions_float3['y']
    x[:,i,2] = positions_float3['z']

    t0 = time.time()
    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), np.int32(1), **gpu_kwargs)
    cuda.Context.synchronize()
    elapsed = time.time() - t0
    print 'elapsed %f sec; %f photons/sec ' % (elapsed, nphotons/elapsed)

    cuda.memcpy_dtoh(positions_float3, positions_gpu)
    cuda.memcpy_dtoh(directions_float3, directions_gpu)
    cuda.memcpy_dtoh(states, states_gpu)
    cuda.memcpy_dtoh(last_hit_triangles, last_hit_triangles_gpu)

    print np.unique(states)

from chromaticity import map_wavelength

#mask = states != 0
mask = geometry.solid_id[last_hit_triangles] == 1

print 'mask length = ', len(states[mask])

rgb_colors = map_wavelength(wavelengths[mask])
colors = rgb_colors*255
colors = colors.astype(np.uint32)
colors = colors[:,0] << 16 | colors[:,1] << 8 | colors[:,2]

def format_hex_string(s):
    return '#' + s.rstrip('L')[2:].zfill(6)

colors = map(format_hex_string, map(hex, colors))

plt.figure()
plt.plot(*roundrobin(x[mask,:,0], x[mask,:,1], colors))
plt.show()

size = (600, 600)

y = np.linspace(np.min(film_mesh.vertices[:,1]), np.max(film_mesh.vertices[:,1]), size[0])
z = np.linspace(np.min(film_mesh.vertices[:,2]), np.max(film_mesh.vertices[:,2]), size[1])

film_positions = np.empty((len(positions_float3[mask]), 3))
film_positions[:,0] = positions_float3[mask]['x']
film_positions[:,1] = positions_float3[mask]['y']
film_positions[:,2] = positions_float3[mask]['z']

film_positions -= film_solid.displacement
film_positions = np.inner(film_positions, np.linalg.inv(film_solid.rotation))

pixels = np.zeros((size[0], size[1], 3))
for position, rgb_color in zip(film_positions, rgb_colors):
    ybin = np.digitize(np.array([position[1]]), y)[0]
    zbin = np.digitize(np.array([position[2]]), z)[0]

    try:
        pixels[ybin,zbin] += rgb_color
    except IndexError:
        continue

import pygame

pixels = ((pixels/np.max(pixels))*255).astype(np.uint32)
pixels = pixels[:,:,0] << 16 | pixels[:,:,1] << 8 | pixels[:,:,2]

pygame.init()
screen = pygame.display.set_mode(size)
pygame.surfarray.blit_array(screen, pixels)
pygame.display.flip()

raw_input()