#!/usr/bin/env python
import time
import os
import numpy as np

from chroma import generator
from chroma import gpu
from chroma import event
from chroma import itertoolset

def pick_seed():
    """Returns a seed for a random number generator selected using
    a mixture of the current time and the current process ID."""
    return int(time.time()) ^ (os.getpid() << 16)

class Simulation(object):
    def __init__(self, detector, seed=None, cuda_device=None,
                 geant4_processes=4, nthreads_per_block=64, max_blocks=1024):
        self.detector = detector

        self.nthreads_per_block = nthreads_per_block
        self.max_blocks = max_blocks

        if seed is None:
            self.seed = pick_seed()
        else:
            self.seed = seed

        # We have three generators to seed: numpy.random, GEANT4, and CURAND.
        # The latter two are done below.
        np.random.seed(self.seed)

        if geant4_processes > 0:
            self.photon_generator = generator.photon.G4ParallelGenerator(geant4_processes, detector.detector_material, base_seed=self.seed)
        else:
            self.photon_generator = None

        self.context = gpu.create_cuda_context(cuda_device)

        if not hasattr(detector, 'mesh'):
            # need to build geometry
            detector.build()

        self.gpu_geometry = gpu.GPUGeometry(detector)
        self.gpu_daq = gpu.GPUDaq(self.gpu_geometry, max(self.detector.pmtids))
        self.gpu_pdf = gpu.GPUPDF()

        self.rng_states = gpu.get_rng_states(self.nthreads_per_block*self.max_blocks, seed=self.seed)

        self.pdf_config = None

    def simulate(self, iterable, keep_photons_beg=False,
                 keep_photons_end=False, run_daq=True, max_steps=10):
        try:
            first_element, iterable = itertoolset.peek(iterable)
        except TypeError:
            first_element, iterable = iterable, [iterable]

        if isinstance(first_element, event.Event):
            iterable = self.photon_generator.generate_events(iterable)
        elif isinstance(first_element, event.Photons):
            iterable = (event.Event(photons_beg=x) for x in iterable)
        elif isinstance(first_element, event.Vertex):
            iterable = (event.Event(primary_vertex=vertex, vertices=[vertex]) for vertex in iterable)
            iterable = self.photon_generator.generate_events(iterable)

        for ev in iterable:
            gpu_photons = gpu.GPUPhotons(ev.photons_beg)

            gpu_photons.propagate(self.gpu_geometry, self.rng_states,
                                  nthreads_per_block=self.nthreads_per_block,
                                  max_blocks=self.max_blocks,
                                  max_steps=max_steps)

            ev.nphotons = len(ev.photons_beg.pos)

            if not keep_photons_beg:
                ev.photons_beg = None

            if keep_photons_end:
                ev.photons_end = gpu_photons.get()

            if run_daq:
                gpu_channels = self.gpu_daq.acquire(gpu_photons, self.rng_states, nthreads_per_block=self.nthreads_per_block, max_blocks=self.max_blocks)
                ev.channels = gpu_channels.get()

            yield ev

    def create_pdf(self, iterable, tbins, trange, qbins, qrange, nreps=1):
        """Returns tuple: 1D array of channel hit counts, 3D array of
        (channel, time, charge) pdfs."""
        first_element, iterable = itertoolset.peek(iterable)

        if isinstance(first_element, event.Event):
            iterable = self.photon_generator.generate_events(iterable)

        pdf_config = (tbins, trange, qbins, qrange)
        if pdf_config != self.pdf_config:
            self.pdf_config = pdf_config
            self.gpu_pdf.setup_pdf(max(self.detector.pmtids), tbins, trange,
                                   qbins, qrange)
        else:
            self.gpu_pdf.clear_pdf()

        if nreps > 1:
            iterable = itertoolset.repeating_iterator(iterable, nreps)

        for ev in iterable:
            gpu_photons = gpu.GPUPhotons(ev.photons_beg)
            gpu_photons.propagate(self.gpu_geometry, self.rng_states,
                                  nthreads_per_block=self.nthreads_per_block,
                                  max_blocks=self.max_blocks)
            gpu_channels = self.gpu_daq.acquire(gpu_photons, self.rng_states, nthreads_per_block=self.nthreads_per_block, max_blocks=self.max_blocks)
            self.gpu_pdf.add_hits_to_pdf(gpu_channels)
        
        return self.gpu_pdf.get_pdfs()

    def eval_pdf(self, event_channels, iterable, min_twidth, trange, min_qwidth, qrange, min_bin_content=20, nreps=1, ndaq=1, time_only=True):
        """Returns tuple: 1D array of channel hit counts, 1D array of PDF
        probability densities."""
        self.gpu_pdf.setup_pdf_eval(event_channels.hit,
                                    event_channels.t,
                                    event_channels.q,
                                    min_twidth,
                                    trange,
                                    min_qwidth,
                                    qrange,
                                    min_bin_content=min_bin_content,
                                    time_only=True)

        first_element, iterable = itertoolset.peek(iterable)

        if isinstance(first_element, event.Event):
            iterable = self.photon_generator.generate_events(iterable)

        for ev in iterable:
            gpu_photons = gpu.GPUPhotons(ev.photons_beg, ncopies=nreps)
            gpu_photons.propagate(self.gpu_geometry, self.rng_states,
                                  nthreads_per_block=self.nthreads_per_block,
                                  max_blocks=self.max_blocks)
            for gpu_photon_slice in gpu_photons.iterate_copies():
                for idaq in xrange(ndaq):
                    gpu_channels = self.gpu_daq.acquire(gpu_photon_slice, self.rng_states, nthreads_per_block=self.nthreads_per_block, max_blocks=self.max_blocks)
                    self.gpu_pdf.accumulate_pdf_eval(gpu_channels)
        
        return self.gpu_pdf.get_pdf_eval()

    def __del__(self):
        self.context.pop()