import numpy as np
import tables
tables.print_versions()

LEN_PMT = int(1.2e6)
NPMTS = 12
NEVENTS = 10

!rm PMT*.h5

def gaussian(x, mu, sig):
    return np.exp(-np.power(x - mu, 2.) / (2 * np.power(sig, 2.)))

x = np.linspace(0, 1, 1e7)
rd = (gaussian(x, 1, 1.) * 1e6).astype(np.int32)

def raw_data(length):
    # Return the actual data that you think it represents PM waveforms better 
    #return np.arange(length, dtype=np.int32)
    return rd[:length]

"""
Explanation: Using a single Table vs EArray + Table
The PyTables community keep asking what can be considered a FAQ.  Namely, should I use a single Table for storing my data, or should I split it in a Table and an Array?
Although there is not a totally general answer, the study below address this for the common case where one has 'raw data' and other data that can be considered 'meta'.  See for example: https://groups.google.com/forum/#!topic/pytables-users/vBEiaRzp3gI
End of explanation
"""


class PMTRD(tables.IsDescription):
    # event_id = tables.Int32Col(pos=1, indexed=True) 
    event_id = tables.Int32Col(pos=1)
    npmt = tables.Int8Col(pos=2)
    pmtrd = tables.Int32Col(shape=LEN_PMT, pos=3) 

def one_table(filename, filters):
    with tables.open_file("{}-{}-{}.h5".format(filename, filters.complib, filters.complevel), "w", filters=filters) as h5t:
        pmt = h5t.create_table(h5t.root, "pmt", PMTRD, expectedrows=NEVENTS*NPMTS)
        pmtr = pmt.row
        for i in range(NEVENTS):
            for j in range(NPMTS):
                pmtr['event_id'] = i
                pmtr['npmt'] = j
                pmtr['pmtrd'] = raw_data(LEN_PMT)
                pmtr.append()

# Using no compression
%time one_table("PMTs", tables.Filters(complib="zlib", complevel=0))

# Using Zlib (level 5) compression
%time one_table("PMTs", tables.Filters(complib="zlib", complevel=5))

# Using Blosc (level 9) compression
%time one_table("PMTs", tables.Filters(complib="blosc:lz4", complevel=9))

ls -sh *.h5

"""
Explanation: Using Tables to store everything
End of explanation
"""


def rawdata_earray(filename, filters):
    with tables.open_file("{}-{}.h5".format(filename, filters.complib), "w", filters=filters) as h5a:
        pmtrd = h5a.create_earray(h5a.root, "pmtrd", tables.Int32Atom(), shape=(0, NPMTS, LEN_PMT),
                                  chunkshape=(1,1,LEN_PMT))
        for i in range(NEVENTS):
            rdata = []
            for j in range(NPMTS):
                rdata.append(raw_data(LEN_PMT))
            pmtrd.append(np.array(rdata).reshape(1, NPMTS, LEN_PMT))
        pmtrd.flush()

# Using no compression
%time rawdata_earray("PMTAs", tables.Filters(complib="zlib", complevel=0))

# Using Zlib (level 5) compression
%time rawdata_earray("PMTAs", tables.Filters(complib="zlib", complevel=5))

# Using Blosc (level 5) compression
%time rawdata_earray("PMTAs", tables.Filters(complib="blosc:lz4", complevel=9))

!ls -sh *.h5

"""
Explanation: So, using no compression leads to best speed, whereas Zlib can compress data by ~32x.  Zlib is ~3x slower than using no compression though.  On its hand, the Blosc compressor is faster but it can barely compress the dataset.
Using EArrays for storing raw data and Table for other metadata
End of explanation
"""


# Add the event IDs in a separate table in the same file
class PMTRD(tables.IsDescription):
    # event_id = tables.Int32Col(pos=1, indexed=True) 
    event_id = tables.Int32Col(pos=1)
    npmt = tables.Int8Col(pos=2)

def add_table(filename, filters):
    with tables.open_file("{}-{}.h5".format(filename, filters.complib), "a", filters=filters) as h5a:
        pmt = h5a.create_table(h5a.root, "pmt", PMTRD)
        pmtr = pmt.row
        for i in range(NEVENTS):
            for j in range(NPMTS):
                pmtr['event_id'] = i
                pmtr['npmt'] = j
                pmtr.append()

# Using no compression
%time add_table("PMTAs", tables.Filters(complib="zlib", complevel=0))

# Using Zlib (level 5) compression
%time add_table("PMTAs", tables.Filters(complib="zlib", complevel=5))

# Using Blosc (level 9) compression
%time add_table("PMTAs", tables.Filters(complib="blosc:lz4", complevel=9))

!ls -sh *.h5

"""
Explanation: We see that by using the Blosc compressor one can achieve around 10x faster output operation wrt Zlib, although the compression ratio can be somewhat smaller (but still pretty good).
End of explanation
"""


def read_single_table(complib, complevel):
    with tables.open_file("PMTs-{}-{}.h5".format(complib, complevel), "r") as h5t:
        pmt = h5t.root.pmt
        for i, row in enumerate(pmt):
            event_id, npmt, pmtrd = row["event_id"], row["npmt"], row["pmtrd"][:]
            if i % 20 == 0:
                print(event_id, npmt, pmtrd[0:5])

%time read_single_table("None", 0)

%time read_single_table("zlib", 5)

%time read_single_table("blosc:lz4", 9)

"""
Explanation: After adding the table we continue to see that a better compression ratio is achieved for EArray + Table with respect to a single Table.  Also, Blosc can make writing files significantly faster than not using compression (it has to write less).
Retrieving data from a single Table
End of explanation
"""


def read_earray_table(complib, complevel):
    with tables.open_file("PMTAs-{}.h5".format(complib, "r")) as h5a:
        pmt = h5a.root.pmt
        pmtrd_ = h5a.root.pmtrd
        for i, row in enumerate(pmt):
            event_id, npmt = row["event_id"], row["npmt"]
            pmtrd = pmtrd_[event_id, npmt]
            if i % 20 == 0:
                print(event_id, npmt, pmtrd[0:5])

%time read_earray_table("None", 0)

%time read_earray_table("zlib", 5)

%time read_earray_table("blosc:lz4", 9)

"""
Explanation: As Blosc could not compress the table, it has a performance that is worse (quite worse actually) to the uncompressed table.  On its hand, Zlib can be more than 3x slower for reading than without compression.
Retrieving data from the EArray + Table
End of explanation
"""


import numpy as np
import matplotlib
import matplotlib.pyplot as plt
%matplotlib inline

"""
Explanation: So, the EArray + Table takes a similar time to read than a pure Table approach when no compression is used. And for some reason, when Zlib is used for compressing the data, the EArray + Table scenario degrades read speed significantly.  However, when the Blosc compression is used, the EArray + Table works actually faster than for the single Table. 
Some plots on speeds and sizes
End of explanation
"""


fig, (ax1, ax2) = plt.subplots(nrows=2, sharex=False, sharey=False)
fig.set_size_inches(w=8, h=7, forward=False)
f = .550   # conversion factor to GB/s
rects1 = ax1.bar(np.arange(3), f / np.array([.323, 3.17, .599]), 0.25, color='r')
rects2 = ax1.bar(np.arange(3) + 0.25, f / np.array([.250, 2.74, .258]), 0.25, color='y')
_ = ax1.set_ylabel('GB/s')
_ = ax1.set_xticks(np.arange(3) + 0.25)
_ = ax1.set_xticklabels(('No Compr', 'Zlib', 'Blosc'))
_ = ax1.legend((rects1[0], rects2[0]), ('Single Table', 'EArray+Table'), loc=9)
_ = ax1.set_title('Speed to store data')

rects1 = ax2.bar(np.arange(3), f / np.array([.099, .592, .782]), 0.25, color='r')
rects2 = ax2.bar(np.arange(3) + 0.25, f / np.array([.082, 1.09, .171]), 0.25, color='y')
_ = ax2.set_ylabel('GB/s')
_ = ax2.set_xticks(np.arange(3) + 0.25)
_ = ax2.set_xticklabels(('No Compr', 'Zlib', 'Blosc'))
_ = ax2.legend((rects1[0], rects2[0]), ('Single Table', 'EArray+Table'), loc=9)
_ = ax2.set_title('Speed to read data')

"""
Explanation: Let's have a look at the speeds at which data can be stored and read using the different paradigms:
End of explanation
"""


fig, ax1 = plt.subplots()
fig.set_size_inches(w=8, h=5)
rects1 = ax1.bar(np.arange(3), np.array([550, 17, 42]), 0.25, color='r')
rects2 = ax1.bar(np.arange(3) + 0.25, np.array([550, 4.1, 9]), 0.25, color='y')
_ = ax1.set_ylabel('MB')
_ = ax1.set_xticks(np.arange(3) + 0.25)
_ = ax1.set_xticklabels(('No Compr', 'Zlib', 'Blosc'))
_ = ax1.legend((rects1[0], rects2[0]), ('Single Table', 'EArray+Table'), loc=9)
_ = ax1.set_title('Size for stored datasets')

"""
Explanation: And now, see the different sizes for the final files:
End of explanation
"""


