#!/usr/bin/env python3

import os
import queue
import threading

import numpy as np
import tables as tb


SIZE = 100
NTHREADS = 5
FILENAME = 'simple_threading.h5'
H5PATH = '/array'


def create_test_file(filename):
    data = np.random.rand(SIZE, SIZE)

    with tb.open_file(filename, 'w') as h5file:
        h5file.create_array('/', 'array', title="Test Array", obj=data)


def chunk_generator(data_size, nchunks):
    chunk_size = int(np.ceil(data_size / nchunks))
    for start in range(0, data_size, chunk_size):
        yield slice(start, start + chunk_size)


lock = threading.Lock()


def synchronized_open_file(*args, **kwargs):
    with lock:
        return tb.open_file(*args, **kwargs)


def synchronized_close_file(self, *args, **kwargs):
    with lock:
        return self.close(*args, **kwargs)


def run(filename, path, inqueue, outqueue):
    try:
        yslice = inqueue.get()
        h5file = synchronized_open_file(filename, mode='r')
        h5array = h5file.get_node(path)
        data = h5array[yslice, ...]
        psum = np.sum(data)
    except Exception as e:
        outqueue.put(e)
    else:
        outqueue.put(psum)
    finally:
        synchronized_close_file(h5file)


def main():
    # generate the test data
    if not os.path.exists(FILENAME):
        create_test_file(FILENAME)

    threads = []
    inqueue = queue.Queue()
    outqueue = queue.Queue()

    # start all threads
    for i in range(NTHREADS):
        thread = threading.Thread(target=run,
                                  args=(FILENAME, H5PATH, inqueue, outqueue))
        thread.start()
        threads.append(thread)

    # push requests in the input queue
    for yslice in chunk_generator(SIZE, len(threads)):
        inqueue.put(yslice)

    # collect results
    try:
        mean_ = 0.

        for i in range(len(threads)):
            out = outqueue.get()
            if isinstance(out, Exception):
                raise out
            else:
                mean_ += out

        mean_ /= SIZE * SIZE

    finally:
        for thread in threads:
            thread.join()

    # print results
    print('Mean: {}'.format(mean_))


if __name__ == '__main__':
    main()