#!/usr/bin/env python
#--*-python-*-

if __name__ == '__main__':
   import optparse
   import sys
   import imp
   import os
   import inspect
   import numpy as np

   from chroma import event
   from chroma import itertoolset
   from chroma import Simulation
   from chroma import root

   parser = optparse.OptionParser('%prog <detector>')
   parser.add_option('-o', dest='output_filename',
                     help='output filename', default='out.root')
   parser.add_option('-j', type='int', dest='device',
                     help='CUDA device number', default=None)
   parser.add_option('-s', type='int', dest='seed',
                     help='random number generator seed')
   parser.add_option('-g', type='int', dest='ngenerators',
                     help='number of GEANT4 processes', default=4)
   parser.add_option('-n', '--nevents', type='int', dest='nevents',
                     default=100)
   parser.add_option('-p', '--particle', dest='particle',
                     help='particle name', default='e-')
   parser.add_option('-k', '--ke', type='float', dest='ke',
                     help='kinetic energy (MeV)', default=100.0)
   parser.add_option('--pos', dest='pos',
                     help='particle vertex origin', default='0,0,0')
   parser.add_option('--dir', dest='dir',
                     help='particle vertex direction', default='1,0,0')
   parser.add_option('--save-photons-beg', action='store_true',
                     dest='save_photons_beg',
                     help='save initial photons to disk', default=False)
   parser.add_option('--save-photons-end', action='store_true',
                     dest='save_photons_end',
                     help='save final photons to disk', default=False)

   options, args = parser.parse_args()

   if len(args) < 1:
      sys.exit(parser.format_help())

   name, attr = args[0].split('.')

   try:
      file, path, description = imp.find_module(name)
   except ImportError:
      raise

   module = imp.load_module(name, file, path, description)

   detector = getattr(module, attr)

   if inspect.isfunction(detector):
      detector = detector()

   pos = np.array([float(s) for s in options.pos.split(',')], dtype=float)
   dir = np.array([float(s) for s in options.dir.split(',')], dtype=float)

   vertex = event.Vertex(options.particle, pos, dir, options.ke)

   sim = Simulation(detector, seed=options.seed, cuda_device=options.device,
                    geant4_processes=options.ngenerators)

   print 'RNG seed: %i' % sim.seed

   writer = root.RootWriter(options.output_filename)

   for i, ev in \
          enumerate(sim.simulate(itertoolset.repeatcopy(vertex,
                                                        options.nevents),
                                 keep_photons_beg=options.save_photons_beg,
                                 keep_photons_end=options.save_photons_end)):
      print '\rEvent: %i' % (i+1),
      sys.stdout.flush()

      assert ev.nphotons > 0, 'Geant4 generated event with no photons!'

      writer.write_event(ev)
   print

   writer.close()