import os
import numpy as np
import pickle
from chroma import *
from copy import deepcopy
from histogram import *

models_directory = os.path.split(os.path.realpath(__file__))[0] + '/../models'

strings = 20
pmts_per_string = 10
radius = 10.0
height = 20.0

grid_spacing = height/pmts_per_string

block_size = 64


class LBNE(geometry.Geometry):
    def __init__(self):
        super(LBNE, self).__init__()

        pmt_mesh = stl.read_stl(models_directory + '/hamamatsu_12inch.stl')
        pmt_mesh /= 1000.0

        apmt = geometry.Solid(pmt_mesh, materials.glass, materials.h2o)

        self.pmt_index = []
        self.pmt_local_axes = []
        self.pmt_positions = []

        self.pmt_hits = []

        for i in range(pmts_per_string):
            for j in range(strings):
                pmt = deepcopy(apmt)
                pmt.mesh += (-radius,0,i*(height/pmts_per_string))
                pmt.mesh = transform.rotate(pmt.mesh, j*2*np.pi/strings, (0,0,1))
                self.add_solid(pmt)
                self.pmt_hits.append(Histogram(10000, (-0.5, 9999.5)))

        for x in np.arange(-radius, radius, grid_spacing):
            for y in np.arange(-radius, radius, grid_spacing):
                if np.sqrt(x**2+y**2) <= radius:
                    pmt = deepcopy(apmt)
                    pmt.mesh = transform.rotate(pmt.mesh, np.pi/2, (0,1,0))
                    pmt.mesh += (x,y,0)
                    self.add_solid(pmt)
                    self.pmt_hits.append(Histogram(10000, (-0.5, 9999.5)))

        for x in np.arange(-radius, radius, grid_spacing):
            for y in np.arange(-radius, radius, grid_spacing):
                if np.sqrt(x**2+y**2) <= radius:
                    pmt = deepcopy(apmt)
                    pmt.mesh = transform.rotate(pmt.mesh, -np.pi/2, (0,1,0))
                    pmt.mesh += (x,y,height)
                    self.add_solid(pmt)
                    self.pmt_hits.append(Histogram(10000, (-0.5, 9999.5)))

        self.build(bits=4)

        self.npmts = len(self.pmt_hits)

        self.gpu = gpu.GPU()
        self.gpu.load_geometry(self)

    def throw_photon_bomb(self, z_position, nphotons=100000):
        origin = np.zeros((nphotons,3)) + (0,0,z_position)

        direction = photon.uniform_sphere(nphotons)

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

        pixels = np.empty(nphotons, dtype=np.int32)
        states = np.empty(nphotons, dtype=np.int32)

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

        gpu_kwargs = {'block': (block_size,1,1), 'grid': (nphotons//block_size+1,1)}
        self.gpu.call(np.int32(nphotons), origin_gpu, direction_gpu, np.int32(self.first_leaf), states_gpu, pixels_gpu, **gpu_kwargs)

        gpu.cuda.memcpy_dtoh(states, states_gpu)

        pmt_indices = self.solid_index[states[(states != -1)]]

        bin_count = np.bincount(pmt_indices)

        bin_count = np.append(bin_count, np.zeros(self.npmts-bin_count.size))

        return bin_count

    def generate_event(self, z_position):
        self.bin_count = self.throw_photon_bomb(z_position)

    def get_likelihood(self, z_position, calls=1000):
        if not hasattr(self, 'bin_count'):
            raise Exception('must call generate_event() first')

        for pmt_hit in self.pmt_hits:
            pmt_hit.reset()

        for i in range(calls):
            print 'throwing bomb %i' % i
            bin_count = self.throw_photon_bomb(z_position)

            for i, count in enumerate(bin_count):
                self.pmt_hits[i].fill(count)

        for pmt_hit in self.pmt_hits:
            pmt_hit.normalize()

        likelihood = 0.0
        for i in range(self.npmts):
            probability = self.pmt_hits[i].eval(self.bin_count[i])

            if probability == 0.0:
                print 'calculating likelihood from pmt %i' % i
                print 'bin count =', self.bin_count[i]
                print self.pmt_hits[i].hist

            likelihood -= np.log(self.pmt_hits[i].eval(self.bin_count[i]))

        return likelihood

if __name__ == '__main__':
    lbne = LBNE()
    view.view(lbne)