import sys
import time
import numpy as np
from pycuda import autoinit
from pycuda.compiler import SourceModule
from pycuda import gpuarray
import pycuda.driver as cuda
from uncertainties import ufloat, umath
from detectors import LBNE
from histogram import Histogram
from photon import uniform_sphere
import layout

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

source = open(layout.source + '/kernel.cu').read()
module = SourceModule(source, options=['-I' + layout.source], no_extern_c=True, cache_dir=False)
propagate = module.get_function('propagate')

lbne = LBNE()
lbne.build(bits=9)
texrefs = lbne.load(module)

nblocks = 64

def generate_event(z=0.0, nphotons=1000):
    origins = np.tile(gpuarray.vec.make_float3(0,0,z), (nphotons,1))
    origins_gpu = cuda.to_device(origins)

    directions = uniform_sphere(nphotons)
    directions_float3 = np.empty(nphotons, 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)

    dest = np.empty(nphotons, dtype=np.int32)
    dest_gpu = cuda.to_device(dest)

    propagate(np.int32(nphotons), origins_gpu, directions_gpu, np.int32(lbne.node_map.size-1), np.int32(lbne.first_node), dest_gpu, block=(nblocks,1,1), grid=(nphotons//nblocks+1,1), texrefs=texrefs)
    cuda.Context.synchronize()

    cuda.memcpy_dtoh(dest, dest_gpu)

    triangles = dest[(dest != -1)]

    event_bincount = np.zeros(len(lbne.solids))
    gpu_bincount = np.bincount(lbne.solid_index[triangles])
    event_bincount[:gpu_bincount.size] = gpu_bincount

    return event_bincount

def likelihood(event_bincount, z=0.0, nphotons=1000, neval=1000):
    origins = np.tile(gpuarray.vec.make_float3(0,0,z), (nphotons,1))
    origins_gpu = cuda.to_device(origins)

    bincount = np.zeros((neval, len(lbne.solids)))
    for i in range(neval):
        print '\revent: %i' % (i+1),
        sys.stdout.flush()

        directions = uniform_sphere(nphotons)
        directions_float3 = np.empty(nphotons, 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)

        dest = np.empty(nphotons, dtype=np.int32)
        dest_gpu = cuda.to_device(dest)

        propagate(np.int32(nphotons), origins_gpu, directions_gpu, np.int32(lbne.node_map.size-1), np.int32(lbne.first_node), dest_gpu, block=(nblocks,1,1), grid=(nphotons//nblocks+1,1))
        cuda.Context.synchronize()

        cuda.memcpy_dtoh(dest, dest_gpu)

        triangles = dest[(dest != -1)]

        gpu_bincount = np.bincount(lbne.solid_index[triangles])
        bincount[i][:gpu_bincount.size] = gpu_bincount
    print

    log_likelihood = ufloat((0,0))
    for i in range(len(lbne.solids)):
        h = Histogram(100, (-0.5, 99.5))
        h.fill(bincount[:,i])
        h.normalize()

        probability = h.ueval(event_bincount[i])

        if probability.nominal_value == 0.0:
            probability = ufloat((0.5/h.nentries, 0.5/h.nentries))

        log_likelihood += umath.log(probability)

    return -log_likelihood