from tables import *
import numpy

'''
The way these "IsDescription constructor" work is simple: write the
code as if it were in a file, then exec()ute it, leaving us with
a local-scoped LocalDescription which may be used to call createTable.

It's a small hack, but it's necessary as the names of the columns
are retrieved based on the variable name, which we can't programmatically set
otherwise.
'''

def get_beginning_description_n_ints(int_names, int_width=64):
    """
    Begins construction of a class inheriting from IsDescription
    to construct an HDF5 table with index columns named with int_names.

    See Series().__init__ to see how those are used.
    """
    int_constructor = "Int64Col"
    if int_width == 32:
        int_constructor = "Int32Col"

    toexec = "class LocalDescription(IsDescription):\n"

    pos = 0

    for n in int_names:
        toexec += "\t" + n + " = " + int_constructor + "(pos=" + str(pos) + ")\n"

    return toexec

def get_description_with_n_ints_n_floats(int_names, float_names, int_width=64, float_width=32):
    """
    Constructs a class to be used when constructing a table with PyTables.

    This is useful to construct a series with an index with multiple levels.
    E.g. if you want to index your "validation error" with "epoch" first, then
    "minibatch_index" second, you'd use two "int_names".

    Parameters
    ----------
    int_names : tuple of str
        Names of the int (e.g. index) columns
    float_names : tuple of str
        Names of the float (e.g. error) columns
    int_width : {'32', '64'}
        Type of ints.
    float_width : {'32', '64'}
        Type of floats.

    Returns
    -------
    A class object, to pass to createTable()
    """

    toexec = get_beginning_description_n_ints(int_names, int_width=int_width)

    float_constructor = "Float32Col"
    if float_width == 64:
        float_constructor = "Float64Col"
    
    pos = len(int_names)

    for n in float_names:
        toexec += "\t" + n + " = " + float_constructor + "(pos=" + str(pos) + ")\n"

    exec(toexec)

    return LocalDescription

class Series():
    def __init__(self, table_name, hdf5_file, index_names=('epoch',), title=None, hdf5_group='/'):
        """Basic arguments each Series must get.

        Parameters
        ----------
        table_name : str
            Name of the table to create under group "hd5_group" (other parameter). No spaces, ie. follow variable naming restrictions.
        hdf5_file : open HDF5 file
            File opened with openFile() in PyTables (ie. return value of openFile).
        index_names : tuple of str
            Columns to use as index for elements in the series, other example would be ('epoch', 'minibatch'). This would then allow you to call append(index, element) with index made of two ints, one for epoch index, one for minibatch index in epoch.
        title : str
            Title to attach to this table as metadata. Can contain spaces and be longer then the table_name.
        hdf5_group : str
            Path of the group (kind of a file) in the HDF5 file under which to create the table.
        """
        self.table_name = table_name
        self.hdf5_file = hdf5_file
        self.index_names = index_names
        self.title = title

    def append(self, index, element):
        raise NotImplementedError

# To put in a series dictionary instead of a real series, to do nothing
# when we don't want a given series to be saved.
class DummySeries():
    def append(self, index, element):
        pass

class ErrorSeries(Series):
    def __init__(self, error_name, table_name, hdf5_file, index_names=('epoch',), title=None, hdf5_group='/'):
        Series.__init__(self, table_name, hdf5_file, index_names, title)

        self.error_name = error_name

        table_description = self._get_table_description()

        self._table = hdf5_file.createTable(hdf5_group, self.table_name, table_description, title=title)

    def _get_table_description(self):
        return get_description_with_n_ints_n_floats(self.index_names, (self.error_name,))

    def append(self, index, error):
        """
        Parameters
        ----------
        index : tuple of int
            Following index_names passed to __init__, e.g. (12, 15) if index_names were ('epoch', 'minibatch_size')
        error : float
            Next error in the series.
        """
        if len(index) != len(self.index_names):
            raise ValueError("index provided does not have the right length (expected " \
                            + str(len(self.index_names)) + " got " + str(len(index)))

        newrow = self._table.row

        # Columns for index in table are based on index_names
        for col_name, value in zip(self.index_names, index):
            newrow[col_name] = value
        newrow[self.error_name] = error

        newrow.append()

        self.hdf5_file.flush()

# Does not inherit from Series because it does not itself need to
# access the hdf5_file and does not need a series_name (provided
# by the base_series.)
class AccumulatorSeriesWrapper():
    """
    
    """
    def __init__(self, base_series, reduce_every, reduce_function=numpy.mean):
        """
        Parameters
        ----------
        base_series : Series
            This object must have an append(index, value) function.
        reduce_every : int
            Apply the reduction function (e.g. mean()) every time we get this number of elements. E.g. if this is 100, then every 100 numbers passed to append(), we'll take the mean and call append(this_mean) on the BaseSeries.
        reduce_function : function
            Must take as input an array of "elements", as passed to (this accumulator's) append(). Basic case would be to take an array of floats and sum them into one float, for example.
        """
        self.base_series = base_series
        self.reduce_function = reduce_function
        self.reduce_every = reduce_every

        self._buffer = []

    
    def append(self, index, element):
        """
        Parameters
        ----------
        index : tuple of int
            The index used is the one of the last element reduced. E.g. if
            you accumulate over the first 1000 minibatches, the index
            passed to the base_series.append() function will be 1000.
        element : float
            Element that will be accumulated.
        """
        self._buffer.append(element)

        if len(self._buffer) == self.reduce_every:
            reduced = self.reduce_function(self._buffer)
            self.base_series.append(index, reduced)
            self._buffer = []

        # This should never happen, except if lists
        # were appended, which should be a red flag.
        assert len(self._buffer) < self.reduce_every

# Outside of class to fix an issue with exec in Python 2.6.
# My sorries to the God of pretty code.
def _BasicStatisticsSeries_construct_table_toexec(index_names):
    toexec = get_beginning_description_n_ints(index_names)

    bpos = len(index_names)
    toexec += "\tmean = Float32Col(pos=" + str(bpos) + ")\n"
    toexec += "\tmin = Float32Col(pos=" + str(bpos+1) + ")\n"
    toexec += "\tmax = Float32Col(pos=" + str(bpos+2) + ")\n"
    toexec += "\tstd = Float32Col(pos=" + str(bpos+3) + ")\n"
    
    # This creates "LocalDescription", which we may then use
    exec(toexec)

    return LocalDescription

basic_stats_functions = {'mean': lambda(x): numpy.mean(x),
                    'min': lambda(x): numpy.min(x),
                    'max': lambda(x): numpy.max(x),
                    'std': lambda(x): numpy.std(x)}

class BasicStatisticsSeries(Series):
    """
    Parameters
    ----------
    series_name : str
        Not optional here. Will be prepended with "Basic statistics for "
    stats_functions : dict, optional
        Dictionary with a function for each key "mean", "min", "max", "std". The function must take whatever is passed to append(...) and return a single number (float).
    """
    def __init__(self, table_name, hdf5_file, stats_functions=basic_stats_functions, index_names=('epoch',), title=None, hdf5_group='/'):
        Series.__init__(self, table_name, hdf5_file, index_names, title)

        self.hdf5_group = hdf5_group

        self.stats_functions = stats_functions

        self._construct_table()

    def _construct_table(self):
        table_description = _BasicStatisticsSeries_construct_table_toexec(self.index_names)

        self._table = self.hdf5_file.createTable(self.hdf5_group, self.table_name, table_description)

    def append(self, index, array):
        """
        Parameters
        ----------
        index : tuple of int
            Following index_names passed to __init__, e.g. (12, 15) if index_names were ('epoch', 'minibatch_size')
        array
            Is of whatever type the stats_functions passed to __init__ can take. Default is anything numpy.mean(), min(), max(), std() can take. 
        """
        if len(index) != len(self.index_names):
            raise ValueError("index provided does not have the right length (expected " \
                            + str(len(self.index_names)) + " got " + str(len(index)))

        newrow = self._table.row

        for col_name, value in zip(self.index_names, index):
            newrow[col_name] = value

        newrow["mean"] = self.stats_functions['mean'](array)
        newrow["min"] = self.stats_functions['min'](array)
        newrow["max"] = self.stats_functions['max'](array)
        newrow["std"] = self.stats_functions['std'](array)

        newrow.append()

        self.hdf5_file.flush()

class SeriesArrayWrapper():
    """
    Simply redistributes any number of elements to sub-series to respective append()s.

    To use if you have many elements to append in similar series, e.g. if you have an array containing [train_error, valid_error, test_error], and 3 corresponding series, this allows you to simply pass this array of 3 values to append() instead of passing each element to each individual series in turn.
    """

    def __init__(self, base_series_list):
        self.base_series_list = base_series_list

    def append(self, index, elements):
        if len(elements) != len(self.base_series_list):
            raise ValueError("not enough or too much elements provided (expected " \
                            + str(len(self.base_series_list)) + " got " + str(len(elements)))

        for series, el in zip(self.base_series_list, elements):
            series.append(index, el)

class SharedParamsStatisticsWrapper(SeriesArrayWrapper):
    '''Save mean, min/max, std of shared parameters place in an array.

    This is specifically for cases where we have _shared_ parameters,
    as we take the .value of each array'''

    def __init__(self, arrays_names, new_group_name, hdf5_file, base_group='/', index_names=('epoch',), title=""):
        """
        Parameters
        ----------
        array_names : array of str
            Name of each array, in order of the array passed to append(). E.g. ('layer1_b', 'layer1_W', 'layer2_b', 'layer2_W')
        new_group_name : str
            Name of a new HDF5 group which will be created under base_group to store the new series.
        base_group : str
            Path of the group under which to create the new group which will store the series.
        title : str
            Here the title is attached to the new group, not a table.
        """
        base_series_list = []

        new_group = hdf5_file.createGroup(base_group, new_group_name, title=title)

        stats_functions = {'mean': lambda(x): numpy.mean(x.value),
                    'min': lambda(x): numpy.min(x.value),
                    'max': lambda(x): numpy.max(x.value),
                    'std': lambda(x): numpy.std(x.value)}

        for name in arrays_names:
            base_series_list.append(
                        BasicStatisticsSeries(
                                table_name=name,
                                hdf5_file=hdf5_file,
                                index_names=index_names,
                                stats_functions=stats_functions,
                                hdf5_group=new_group._v_pathname))

        SeriesArrayWrapper.__init__(self, base_series_list)