ALLEN INSTITUTE

# Source code for bmtk.utils.reports.current_writer

    # Copyright 2017. Allen Institute. All rights reserved
    #
    # Redistribution and use in source and binary forms, with or without modification, are permitted provided that the
    # following conditions are met:
    #
    # 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following
    # disclaimer.
    #
    # 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following
    # disclaimer in the documentation and/or other materials provided with the distribution.
    #
    # 3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote
    # products derived from this software without specific prior written permission.
    #
    # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
    # INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
    # DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
    # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
    # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
    # WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
    # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
    #
    import os
    import h5py
    import numpy as np

    from bmtk.utils.sonata.utils import add_hdf5_magic, add_hdf5_version

    class CurrentWriterv01(object):
        """Used to save current currents to the described hdf5 format.
        """

        class DataTable(object):
            """A small struct to keep track of different data (and buffer) tables"""
            def __init__(self):
                # If buffering data, buffer_block will be an in-memory array and will write to data_block during when
                # filled. If not buffering buffer_block is an hdf5 dataset and data_block is ignored
                self.data_block = None
                self.buffer_block = None

        def __init__(self, file_name, num_currents, units=None, tstart=0.0, tstop=1.0, dt=0.01, n_steps=None,
                     buffer_size=0, buffer_data=True, **kwargs):
            self._units = units
            self._tstart = tstart
            self._tstop = tstop
            self._dt = dt

            self._num_currents = num_currents

            self._n_steps = n_steps

            self._buffer_size = buffer_size
            self._buffer_data = buffer_data and buffer_size>0
            self._data_block = self.DataTable()
            self._last_save_indx = 0  # for buffering, used to keep track of last timestep data was saved to disk

            self._file_name = file_name
            self._create_h5file()

            self._buffer_block_size = 0
            self._total_steps = 0

            self._is_initialized = False

        def set_units(self, val):
            self._units = val

        def units(self):
            return self._units

        def set_tstart(self, val):
            self._tstart = val

        def tstart(self):
            return self._tstart

        def set_tstop(self, val):
            self._tstop = val

        def tstop(self):
            return self._tstop

        def set_dt(self, val):
            self._dt = val

        def dt(self):
            return self._dt

        def n_steps(self):
            if self._n_steps is None:
                self._n_steps = int((self._tstop - self._tstart) / self._dt)
            return self._n_steps

        def _create_h5file(self):
            fdir = os.path.dirname(os.path.abspath(self._file_name))
            if not os.path.exists(fdir):
                os.mkdir(fdir)
            self.h5_base = h5py.File(self._file_name, 'w')
            add_hdf5_version(self.h5_base)
            add_hdf5_magic(self.h5_base)

        def initialize(self, **kwargs):
            if self._is_initialized:
                return

            n_steps = self.n_steps()
            if n_steps <= 0:
                raise Exception('A non-zero positive integer num-of-steps is required to initialize the current report.'
                                'Please specify report length using the n_steps parameters (or using appropiate tstop,'
                                'tstart, and dt).')

            #self._calc_offset()
            base_grp = self.h5_base

            self._total_steps = n_steps
            self._buffer_block_size = self._buffer_size

            if self._buffer_data:
                # Set up in-memory block to buffer recorded variables before writing to the dataset.
                self._data_block.buffer_block = np.zeros((self._buffer_size, self._num_currents), dtype=float)

                self._data_block.data_block = base_grp.create_dataset('data', shape=(self.n_steps(), self._num_currents),
                                                                      dtype=float, chunks=True)

                if self._units is not None:
                    self._data_block.data_block.attrs['units'] = self._units

            else:
                # Since we are not buffering data, we just write directly to the on-disk dataset.
                self._data_block.buffer_block = base_grp.create_dataset('data', shape=(self.n_steps(), self._n_segments_all),
                                                                   dtype=float, chunks=True)

                if self._units is not None:
                    self._data_block.buffer_block.attrs['units'] = self._units

            self._is_initialized = True

        def record_clamps(self, vals, tstep):
            """Record clamp currents.

            :param vals: list of currents for each clamp
            :param tstep: time step
            """
            self.initialize()
            buffer_block = self._data_block.buffer_block
            update_index = (tstep - self._last_save_indx)
            buffer_block[update_index, :] = vals

        def flush(self):
            """Move data from memory to dataset"""
            if self._buffer_data:
                blk_beg = self._last_save_indx
                blk_end = blk_beg + self._buffer_block_size
                if blk_end > self._total_steps:
                    # Need to handle the case that simulation doesn't end on a block step
                    blk_end = blk_beg + self._total_steps - blk_beg

                block_size = blk_end - blk_beg
                self._last_save_indx += block_size

                self._data_block.data_block[blk_beg:blk_end, :] = self._data_block.buffer_block[:block_size, :]

        def close(self):
            self.h5_base.close()

        def merge(self):
            pass