import unittest
import numpy as np

import chroma.detectors
from chroma.generator.photon import G4ParallelGenerator
from chroma.generator.vertex import constant_particle_gun
from chroma.optics import water_wcsim
from chroma.gpu import GPU
from chroma.sim import Simulation

class TestPDF(unittest.TestCase):
    def setUp(self):
        self.detector = chroma.detectors.find("microlbne")
        self.detector.build()
        self.vertex_gen = constant_particle_gun('e-', (0,0,0), (1,0,0), 10)

    def testGPUPDF(self):
        '''Create a hit count and (q,t) PDF for 10 MeV events in MicroLBNE'''

        g4generator = G4ParallelGenerator(1, water_wcsim)
        gpu = GPU(0)
        gpu.load_geometry(self.detector)
        gpu.setup_propagate()
        gpu.setup_daq(max(self.detector.pmtids))
        gpu.setup_pdf(max(self.detector.pmtids), 100, (-0.5, 999.5), 10, (-0.5, 9.5))
        
        gpu.clear_pdf()

        for ev in g4generator.generate_events(10, self.vertex_gen):
            gpu.load_photons(ev.photon_start)
            gpu.propagate()
            gpu.run_daq()
            gpu.add_hits_to_pdf()

        hitcount, pdf = gpu.get_pdfs()
        self.assertTrue( (hitcount > 0).any() )
        self.assertTrue( (pdf > 0).any() )

        # Consistency checks
        for i, nhits in enumerate(hitcount):
            self.assertEqual(nhits, pdf[i].sum())

    def testSimPDF(self):
        sim = Simulation(self.detector, water_wcsim)
        hitcount, pdf = sim.create_pdf(100, self.vertex_gen,
                                        100, (-0.5, 999.5), 10, (-0.5, 9.5))

        self.assertTrue( (hitcount > 0).any() )
        self.assertTrue( (pdf > 0).any() )

        # Consistency checks
        for i, nhits in enumerate(hitcount):
            self.assertEqual(nhits, pdf[i].sum())