from gputhread import *
from Queue import Queue
from detectors import LBNE
from photon import uniform_sphere
import numpy as np
from pycuda import gpuarray
import pycuda.driver as cuda
from uncertainties import ufloat, umath
from histogram import Histogram

jobs = Queue()
output = Queue()

def generate_event(pos=(0,0,0), nphotons=1000):
    origins_float3 = np.tile(gpuarray.vec.make_float3(*pos), (nphotons,1))
    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]

    jobs.put(Job(origins_float3, directions_float3))

    jobs.join()

    return output.get()

def likelihood(event_bincount, pos=(0,0,0), nphotons=1000, neval=1000):
    for i in range(neval):
        origins_float3 = np.tile(gpuarray.vec.make_float3(*pos), (nphotons,1))
        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]

        jobs.put(Job(origins_float3, directions_float3))

    jobs.join()

    bincount = np.zeros((neval, len(lbne.solids)))
    for i in range(neval):
        bincount[i] = output.get()

    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

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

    parser = optparse.OptionParser('%prog')
    parser.add_option('-b', type='int', dest='nbits', default=8)
    parser.add_option('-j', type='int', dest='ndevices', default=1)
    parser.add_option('-n', type='int', dest='nblocks', default=64)
    options, args = parser.parse_args()

    lbne = LBNE()
    lbne.build(bits=options.nbits)

    cuda.init()

    gputhreads = []
    for i in range(options.ndevices):
        gputhreads.append(GPUThread(i, lbne, jobs, output, options.nblocks))
        gputhreads[-1].start()

    try:
        event_bincount = generate_event()

        for z in np.linspace(-1.0, 1.0, 100):
            t0 = time.time()
            log_likelihood = likelihood(event_bincount, (z,0,0))
            elapsed = time.time() - t0
            print 'z = %5.2f, likelihood = %s, elapsed %.2f sec' % (z, log_likelihood, elapsed)
    finally:
        for gputhread in gputhreads:
            gputhread.stop()

        for gputhread in gputhreads:
            gputhread.join()