from pax import units, configuration, datastructure
# setup hax for Midway
import hax
from hax import cuts
#Select pax version
pax_version = '6.8.0'

hax.__version__
hax.init(experiment='XENON1T',
         raw_data_access_mode='local',
         raw_data_local_path='/project/lgrandi/xenon1t/raw/',
         minitree_paths=[#'/home/mvargasjara/minitrees/v'+pax_version,
                         '/project2/lgrandi/xenon1t/minitrees/pax_v'+pax_version,
                          '/project/lgrandi/xenon1t/minitrees/pax_v' + pax_version
                         ],
         pax_version_policy= pax_version,
         make_minitrees = False
         )

dsets = hax.runs.tags_selection(include=['*sciencerun1'])
dsets = dsets[(dsets.source__type == 'none')]
run_time_s_tmp = pd.to_timedelta(pd.to_datetime(dsets.end)).dt.seconds - pd.to_timedelta(pd.to_datetime(dsets.start)).dt.seconds
run_time_s_tmp[run_time_s_tmp<0] += 24*3600
dsets['run_time_s'] = run_time_s_tmp
# select data range
dsets = dsets[(dsets.start > pd.to_datetime('12/31/2016')) & (dsets.end < pd.to_datetime('02/01/2018'))]
dsets['start_date'] = dsets.start.dt.date
run_names = dsets.name
# get rid of problematic runs
bad_runs = ['171204_1357','171204_1858','171208_0616','171204_0856','171204_1658','171205_0701']
for run in bad_runs:
    run_names = run_names[run_names != run]
    dsets = dsets[dsets.name != run]


preselect = ['cs1 > 0', 'z_3d_nn <= 0.0', 'cs1<100']

df = hax.minitrees.load(run_names, ['Corrections','Basics','Extended','Fundamentals','Proximity',
                                    'PositionReconstruction'], preselection=preselect, force_reload=False,
                        num_workers = 7)

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-7-a6c400c757de> in <module>()
      3 df = hax.minitrees.load(run_names, ['Corrections','Basics','Extended','Fundamentals','Proximity',
      4                                     'PositionReconstruction'], preselection=preselect, force_reload=False,
----> 5                         num_workers = 7)

/project/lgrandi/anaconda3/envs/pax_head/lib/python3.4/site-packages/hax-2.2.3-py3.4.egg/hax/minitrees.py in load(datasets, treemakers, preselection, force_reload, delayed, num_workers, compute_options, cache_file, remake_cache, event_list)
    509
    510     result = dask.dataframe.from_delayed(
--> 511         partial_results, meta=partial_results[0].compute())
    512
    513     if not delayed:

/project/lgrandi/anaconda3/envs/pax_head/lib/python3.4/site-packages/dask/base.py in compute(self, **kwargs)
     77             Extra keywords to forward to the scheduler ``get`` function.
     78         """
---> 79         return compute(self, **kwargs)[0]
     80
     81     @classmethod

/project/lgrandi/anaconda3/envs/pax_head/lib/python3.4/site-packages/dask/base.py in compute(*args, **kwargs)
    177         dsk = merge(var.dask for var in variables)
    178     keys = [var._keys() for var in variables]
--> 179     results = get(dsk, keys, **kwargs)
    180
    181     results_iter = iter(results)

/project/lgrandi/anaconda3/envs/pax_head/lib/python3.4/site-packages/dask/threaded.py in get(dsk, result, cache, num_workers, **kwargs)
     74     results = get_async(pool.apply_async, len(pool._pool), dsk, result,
     75                         cache=cache, get_id=_thread_get_id,
---> 76                         **kwargs)
     77
     78     # Cleanup pools associated to dead threads

/project/lgrandi/anaconda3/envs/pax_head/lib/python3.4/site-packages/dask/async.py in get_async(apply_async, num_workers, dsk, result, cache, get_id, raise_on_exception, rerun_exceptions_locally, callbacks, dumps, loads, **kwargs)
    491                     _execute_task(task, data)  # Re-execute locally
    492                 else:
--> 493                     raise(remote_exception(res, tb))
    494             state['cache'][key] = res
    495             finish_task(dsk, key, state, results, keyorder.get)

ValueError: No TreeMaker named PositionReconstruction known to hax!

Traceback
---------
  File "/project/lgrandi/anaconda3/envs/pax_head/lib/python3.4/site-packages/dask/async.py", line 268, in execute_task
    result = _execute_task(task, data)
  File "/project/lgrandi/anaconda3/envs/pax_head/lib/python3.4/site-packages/dask/async.py", line 249, in _execute_task
    return func(*args2)
  File "/project/lgrandi/anaconda3/envs/pax_head/lib/python3.4/site-packages/dask/compatibility.py", line 47, in apply
    return func(*args, **kwargs)
  File "/project/lgrandi/anaconda3/envs/pax_head/lib/python3.4/site-packages/hax-2.2.3-py3.4.egg/hax/minitrees.py", line 388, in load_single_dataset
    run_id, treemaker, force_reload=force_reload, event_list=event_list)
  File "/project/lgrandi/anaconda3/envs/pax_head/lib/python3.4/site-packages/hax-2.2.3-py3.4.egg/hax/minitrees.py", line 321, in load_single_minitree
    run_id, treemaker, force_reload=force_reload)
  File "/project/lgrandi/anaconda3/envs/pax_head/lib/python3.4/site-packages/hax-2.2.3-py3.4.egg/hax/minitrees.py", line 186, in check
    treemaker_name, treemaker = get_treemaker_name_and_class(treemaker)
  File "/project/lgrandi/anaconda3/envs/pax_head/lib/python3.4/site-packages/hax-2.2.3-py3.4.egg/hax/minitrees.py", line 607, in get_treemaker_name_and_class
    raise ValueError("No TreeMaker named %s known to hax!" % tm)