import numpy as np
from histogram import HistogramDD
from uncertainties import ufloat, umath
from itertools import izip, compress

class Likelihood(object):
    "Class to evaluate likelihoods for detector events."
    def __init__(self, sim, event=None):
        """
        Args:
            - sim: chroma.sim.Simulation
                The simulation object used to simulate events and build pdfs.
            - event: chroma.event.Event, *optional*
                The detector event being reconstructed. If None, you must call
                set_event() before eval().
        """
        self.sim = sim

        if event is not None:
            self.set_event(event)

        self.pdfs = dict(zip(sim.detector.pmtids,[HistogramDD((100,10), range=[(0.0,500.0),(-0.5,9.5)]) for i in sim.detector.pmtids]))

    def set_event(self, event):
        "Set the detector event being reconstructed."
        self.event = event

    def eval(self, vertex_generator, nevals):
        """
        Return the negative log likelihood that the detector event set in the
        constructor or by set_event() was the result of a particle generated
        by `vertex_generator`.
        """
        for pdf in self.pdfs.values():
            pdf.reset()

        for ev in self.sim.simulate(nevals, vertex_generator):
            for i, t, q in compress(izip(range(ev.channels.hit.size),ev.channels.t,ev.channels.q),ev.channels.hit):
                self.pdfs[i].fill((ev.channels.t[i],ev.channels.q[i]))

        for pdf in self.pdfs.values():
            if pdf.nentries > 0:
                pdf.normalize()

        log_likelihood = ufloat((0,0))

        for i, t, q in compress(izip(range(self.event.channels.hit.size),self.event.channels.t,self.event.channels.q),self.event.channels.hit):
            probability = self.pdfs[i].ueval((t,q)).item()

            if probability.nominal_value == 0.0:
                if self.pdfs[i].nentries > 0:
                    probability = ufloat([0.5/self.pdfs[i].nentries]*2)
                else:
                    probability = ufloat([1.0/self.pdfs[i].hist.size]*2)

            log_likelihood += umath.log(probability)

        return -log_likelihood

if __name__ == '__main__':
    from chroma import detectors
    from chroma.sim import Simulation
    from chroma.optics import water
    from chroma.generator import constant_particle_gun

    detector = detectors.find('lbne')
    sim = Simulation(detector, water)

    event = sim.simulate(1, constant_particle_gun('e-',(0,0,0),(1,0,0),100.0)).next()

    print 'nhit = %i' % np.count_nonzero(event.hits.hit)

    likelihood = Likelihood(sim, event)

    for x in np.linspace(-10.0, 10.0, 10*5+1):
        print 'x = %5.1f, %s' % (x, likelihood.eval(constant_particle_gun('e-',(x,0,0),(1,0,0),100.0),100))