import numpy as np
import time
from pathlib import Path
import pickle

import boost_histogram as bh


n_bins = 100
xmin = -100
xmax = 1e4


#Python implementation to read a cluster file and save the 
#histogram. Can be used to compare with the C++ implementation.

fname = Path('beam_En700eV_-40deg_300V_10us_d0_f0_100.clust')

hist_py = bh.Histogram(bh.axis.Regular(n_bins, xmin, xmax))

t0 = time.perf_counter()
N = 1500 #Maximum frames
with open(fname, 'rb') as f:
    for i in range(N):
        frame_number = np.fromfile(f, count=1, dtype = np.int32)
        if frame_number.size == 0:
            break
        n_clusters = np.fromfile(f, count=1, dtype = np.uint32)[0]
        for i in range(n_clusters):
            xy = np.fromfile(f, count=2, dtype = np.uint16)
            data = np.fromfile(f, count=9, dtype = np.int32)
            hist_py.fill(data)
t = time.perf_counter()-t0

with open(fname.with_suffix('.pkl'), 'wb') as f:
    pickle.dump(hist_py, f)

print(f'Reading {i} frames with python took: {t:.3f}s')