#!/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)
    parser.add_option('--debug', dest='debug', 
                      action='store_true', default=False,
                      help='Start python debugger on exception')
    parser.add_option('--rootlogon', dest='rootlogon',
                      action='store_true', default=False,
                      help='Write a .rootlogon.C file to the current directory for reading Chroma ROOT files.')

    options, args = parser.parse_args()

    if options.rootlogon:
        if os.path.exists('.rootlogon.C'):
            print >>sys.stderr, ".rootlogon.C file already exists!"
            sys.exit(1)
        else:
            with open('.rootlogon.C', 'w') as f:
                print >>f, '{\ngROOT->ProcessLine(".L ~/.chroma/root.C+");\n}'
            print >>sys.stderr, '.rootlogon.C created'
            sys.exit(0)

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

    if options.debug:
        enable_debug_on_crash()

    module_name, function_name = args[0].rsplit('.', 1)

    try:
        module = __import__(module_name, fromlist=[function_name])
    except ImportError:
        raise

    detector = getattr(module, function_name)

    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()