import cProfile

import numpy as np
import sys
import hepfile

# For development
#sys.path.append('../src/hepfile')
#import write as hepfile

def main():
    data = hepfile.initialize()

    hepfile.create_group(data,'jet',counter='njet')
    hepfile.create_dataset(data,['e','px','py','pz'],group='jet',dtype=float)
    hepfile.create_dataset(data,['algorithm'],group='jet',dtype=int)
    hepfile.create_dataset(data,['words'],group='jet',dtype=str)

    hepfile.create_group(data,'muons',counter='nmuon')
    hepfile.create_dataset(data,['e','px','py','pz'],group='muons',dtype=float)

    hepfile.create_dataset(data,['METpx','METpy'],dtype=float)

    event = hepfile.create_single_bucket(data)

    rando_words = ["hi", "bye", "ciao", "aloha"]

    print("Created the datasets and groups...")

#'''
    for i in range(0,100000):

        #hepfile.clear_event(event)

        njet = 17
        event['jet/njet'] = njet

        for n in range(njet):
            event['jet/e'].append(np.random.random())
            event['jet/px'].append(np.random.random())
            event['jet/py'].append(np.random.random())
            event['jet/pz'].append(np.random.random())

            event['jet/algorithm'].append(np.random.randint(-1,1))

            event['jet/words'].append(np.random.choice(rando_words))

        event['METpx'] = np.random.random()
        event['METpy'] = np.random.random()

        #hepfile.pack(data,event,EMPTY_OUT_BUCKET=False)
        return_value = hepfile.pack(data,event,STRICT_CHECKING=True)
        if return_value != 0:
            exit()

    print("Writing the file...")
#hdfile = hepfile.write_to_file('output.hdf5',data)
    hdfile = hepfile.write_to_file('output.hdf5',data,comp_type='gzip',comp_opts=9,verbose=True)
#'''

if __name__ == "__main__":
    profiler = cProfile.Profile()
    profiler.enable()
    main()
    profiler.disable()
    profiler.dump_stats("example.prof")
