from lookup_list import LookupList
Example = LookupList
from misc import unique_elements_list_intersection
from string import join
from sys import maxint
import numpy

class AbstractFunction (Exception): """Derived class must override this function"""
class NotImplementedYet (NotImplementedError): """Work in progress, this should eventually be implemented"""

class DataSet(object):
    """A virtual base class for datasets.

    A DataSet can be seen as a generalization of a matrix, meant to be used in conjunction
    with learning algorithms (for training and testing them): rows/records are called examples, and
    columns/attributes are called fields. The field value for a particular example can be an arbitrary
    python object, which depends on the particular dataset.
    
    We call a DataSet a 'stream' when its length is unbounded (otherwise its __len__ method
    should return sys.maxint).

    A DataSet is a generator of iterators; these iterators can run through the
    examples or the fields in a variety of ways.  A DataSet need not necessarily have a finite
    or known length, so this class can be used to interface to a 'stream' which
    feeds on-line learning (however, as noted below, some operations are not
    feasible or not recommanded on streams).

    To iterate over examples, there are several possibilities:
     - for example in dataset([field1, field2,field3, ...]):
     - for val1,val2,val3 in dataset([field1, field2,field3]):
     - for minibatch in dataset.minibatches([field1, field2, ...],minibatch_size=N):
     - for mini1,mini2,mini3 in dataset.minibatches([field1, field2, field3], minibatch_size=N):
     - for example in dataset::
        print example['x']
     - for x,y,z in dataset:
     Each of these is documented below. All of these iterators are expected
     to provide, in addition to the usual 'next()' method, a 'next_index()' method
     which returns a non-negative integer pointing to the position of the next
     example that will be returned by 'next()' (or of the first example in the
     next minibatch returned). This is important because these iterators
     can wrap around the dataset in order to do multiple passes through it,
     in possibly unregular ways if the minibatch size is not a divisor of the
     dataset length.
    
    To iterate over fields, one can do
     - for field in dataset.fields():
         for field_value in field: # iterate over the values associated to that field for all the dataset examples
     - for field in dataset(field1,field2,...).fields() to select a subset of fields
     - for field in dataset.fields(field1,field2,...) to select a subset of fields
    and each of these fields is iterable over the examples:
     - for field_examples in dataset.fields():
        for example_value in field_examples:
           ...
    but when the dataset is a stream (unbounded length), it is not recommanded to do 
    such things because the underlying dataset may refuse to access the different fields in
    an unsynchronized ways. Hence the fields() method is illegal for streams, by default.
    The result of fields() is a DataSetFields object, which iterates over fields,
    and whose elements are iterable over examples. A DataSetFields object can
    be turned back into a DataSet with its examples() method::
      dataset2 = dataset1.fields().examples()
    and dataset2 should behave exactly like dataset1 (in fact by default dataset2==dataset1).
    
    Note: Fields are not mutually exclusive, i.e. two fields can overlap in their actual content.

    Note: The content of a field can be of any type. Field values can also be 'missing'
    (e.g. to handle semi-supervised learning), and in the case of numeric (numpy array)
    fields (i.e. an ArrayFieldsDataSet), NaN plays the role of a missing value. 
    What about non-numeric values? None.

    Dataset elements can be indexed and sub-datasets (with a subset
    of examples) can be extracted. These operations are not supported
    by default in the case of streams.

     - dataset[:n] returns a dataset with the n first examples.

     - dataset[i1:i2:s] returns a dataset with the examples i1,i1+s,...i2-s.

     - dataset[i] returns an Example.

     - dataset[[i1,i2,...in]] returns a dataset with examples i1,i2,...in.

     - dataset[fieldname] an iterable over the values of the field fieldname across
     the dataset (the iterable is obtained by default by calling valuesVStack
     over the values for individual examples).

     - dataset.<property> returns the value of a property associated with
     the name <property>. The following properties should be supported:
          - 'description': a textual description or name for the dataset
          - 'fieldtypes': a list of types (one per field)

    Datasets can be concatenated either vertically (increasing the length) or
    horizontally (augmenting the set of fields), if they are compatible, using
    the following operations (with the same basic semantics as numpy.hstack
    and numpy.vstack):

     - dataset1 | dataset2 | dataset3 == dataset.hstack([dataset1,dataset2,dataset3])

    creates a new dataset whose list of fields is the concatenation of the list of
    fields of the argument datasets. This only works if they all have the same length.

     - dataset1 & dataset2 & dataset3 == dataset.vstack([dataset1,dataset2,dataset3])

    creates a new dataset that concatenates the examples from the argument datasets
    (and whose length is the sum of the length of the argument datasets). This only
    works if they all have the same fields.

    According to the same logic, and viewing a DataSetFields object associated to
    a DataSet as a kind of transpose of it, fields1 & fields2 concatenates fields of
    a DataSetFields fields1 and fields2, and fields1 | fields2 concatenates their
    examples.

    A dataset can hold arbitrary key-value pairs that may be used to access meta-data
    or other properties of the dataset or associated with the dataset or the result
    of a computation stored in a dataset. These can be accessed through the [key] syntax
    when key is a string (or more specifically, neither an integer, a slice, nor a list).
    
    A DataSet sub-class should always redefine the following methods:
       - __len__ if it is not a stream
       - fieldNames
       - minibatches_nowrap (called by DataSet.minibatches())
       - valuesHStack
       - valuesVStack
    For efficiency of implementation, a sub-class might also want to redefine
       - hasFields
       - __getitem__ may not be feasible with some streams
       - __iter__
    """

    def __init__(self,description=None,fieldtypes=None):
        if description is None:
            # by default return "<DataSetType>(<SuperClass1>,<SuperClass2>,...)"
            description = type(self).__name__ + " ( " + join([x.__name__ for x in type(self).__bases__]) + " )"
        self.description=description
        self.fieldtypes=fieldtypes
    
    class MinibatchToSingleExampleIterator(object):
        """
        Converts the result of minibatch iterator with minibatch_size==1 into
        single-example values in the result. Therefore the result of
        iterating on the dataset itself gives a sequence of single examples
        (whereas the result of iterating over minibatches gives in each
        Example field an iterable object over the individual examples in
        the minibatch).
        """
        def __init__(self, minibatch_iterator):
            self.minibatch_iterator = minibatch_iterator
            self.minibatch = None
        def __iter__(self): #makes for loop work
            return self
        def next(self):
            size1_minibatch = self.minibatch_iterator.next()
            if not self.minibatch:
                self.minibatch = Example(size1_minibatch.keys(),[value[0] for value in size1_minibatch.values()])
            else:
                self.minibatch._values = [value[0] for value in size1_minibatch.values()]
            return self.minibatch
        
        def next_index(self):
            return self.minibatch_iterator.next_index()

    def __iter__(self):
        """Supports the syntax "for i in dataset: ..."

        Using this syntax, "i" will be an Example instance (or equivalent) with
        all the fields of DataSet self.  Every field of "i" will give access to
        a field of a single example.  Fields should be accessible via
        i["fielname"] or i[3] (in the order defined by the elements of the
        Example returned by this iterator), but the derived class is free
        to accept any type of identifier, and add extra functionality to the iterator.

        The default implementation calls the minibatches iterator and extracts the first example of each field.
        """
        return DataSet.MinibatchToSingleExampleIterator(self.minibatches(None, minibatch_size = 1))


    class MinibatchWrapAroundIterator(object):
        """
        An iterator for minibatches that handles the case where we need to wrap around the
        dataset because n_batches*minibatch_size > len(dataset). It is constructed from
        a dataset that provides a minibatch iterator that does not need to handle that problem.
        This class is a utility for dataset subclass writers, so that they do not have to handle
        this issue multiple times, nor check that fieldnames are valid, nor handle the
        empty fieldnames (meaning 'use all the fields').
        """
        def __init__(self,dataset,fieldnames,minibatch_size,n_batches,offset):
            self.dataset=dataset
            self.fieldnames=fieldnames
            self.minibatch_size=minibatch_size
            self.n_batches=n_batches
            self.n_batches_done=0
            self.next_row=offset
            self.L=len(dataset)
            assert offset+minibatch_size<=self.L
            ds_nbatches = (self.L-offset)/minibatch_size
            if n_batches is not None:
                ds_nbatches = max(n_batches,ds_nbatches)
            if fieldnames:
                assert dataset.hasFields(*fieldnames)
            else:
                fieldnames=dataset.fieldNames()
            self.iterator = dataset.minibatches_nowrap(fieldnames,minibatch_size,ds_nbatches,offset)

        def __iter__(self):
            return self

        def next_index(self):
            return self.next_row

        def next(self):
            if self.n_batches and self.n_batches_done==self.n_batches:
                raise StopIteration
            upper = self.next_row+self.minibatch_size
            if upper <=self.L:
                minibatch = self.iterator.next()
            else:
                if not self.n_batches:
                    raise StopIteration
                # we must concatenate (vstack) the bottom and top parts of our minibatch
                # first get the beginning of our minibatch (top of dataset)
                first_part = self.dataset.minibatches_nowrap(fieldnames,self.L-self.next_row,1,self.next_row).next()
                second_part = self.dataset.minibatches_nowrap(fieldnames,upper-self.L,1,0).next()
                minibatch = Example(self.fieldnames,
                                    [self.dataset.valuesVStack(name,[first_part[name],second_part[name]])
                                     for name in self.fieldnames])
            self.next_row=upper
            self.n_batches_done+=1
            if upper >= self.L and self.n_batches:
                self.next_row -= self.L
            return DataSetFields(MinibatchDataSet(minibatch,self.dataset.valuesVStack,
                                                  self.dataset.valuesHStack),
                                 minibatch.keys()))


    minibatches_fieldnames = None
    minibatches_minibatch_size = 1
    minibatches_n_batches = None
    def minibatches(self,
                    fieldnames = minibatches_fieldnames,
                    minibatch_size = minibatches_minibatch_size,
                    n_batches = minibatches_n_batches,
                    offset = 0):
        """
        Return an iterator that supports three forms of syntax:

            for i in dataset.minibatches(None,**kwargs): ...

            for i in dataset.minibatches([f1, f2, f3],**kwargs): ...

            for i1, i2, i3 in dataset.minibatches([f1, f2, f3],**kwargs): ...

        Using the first two syntaxes, "i" will be an indexable object, such as a list,
        tuple, or Example instance. In both cases, i[k] is a list-like container
        of a batch of current examples. In the second case, i[0] is
        list-like container of the f1 field of a batch current examples, i[1] is
        a list-like container of the f2 field, etc.

        Using the first syntax, all the fields will be returned in "i".
        Using the third syntax, i1, i2, i3 will be list-like containers of the
        f1, f2, and f3 fields of a batch of examples on each loop iteration.

        The minibatches iterator is expected to return upon each call to next()
        a DataSetFields object, which is a LookupList (indexed by the field names) whose
        elements are iterable over the minibatch examples, and which keeps a pointer to
        a sub-dataset that can be used to iterate over the individual examples
        in the minibatch. Hence a minibatch can be converted back to a regular
        dataset or its fields can be looked at individually (and possibly iterated over).

        PARAMETERS
        - fieldnames (list of any type, default None):
        The loop variables i1, i2, i3 (in the example above) should contain the
        f1, f2, and f3 fields of the current batch of examples.  If None, the
        derived class can choose a default, e.g. all fields.

        - minibatch_size (integer, default 1)
        On every iteration, the variables i1, i2, i3 will have
        exactly minibatch_size elements. e.g. len(i1) == minibatch_size

        - n_batches (integer, default None)
        The iterator will loop exactly this many times, and then stop.  If None,
        the derived class can choose a default.  If (-1), then the returned
        iterator should support looping indefinitely.

        - offset (integer, default 0)
        The iterator will start at example 'offset' in the dataset, rather than the default.
        
        Note: A list-like container is something like a tuple, list, numpy.ndarray or
        any other object that supports integer indexing and slicing.

        """
        return DataSet.MinibatchWrapAroundIterator(self,fieldnames,minibatch_size,n_batches,offset)

    def minibatches_nowrap(self,fieldnames,minibatch_size,n_batches,offset):
        """
        This is the minibatches iterator generator that sub-classes must define.
        It does not need to worry about wrapping around multiple times across the dataset,
        as this is handled by MinibatchWrapAroundIterator when DataSet.minibatches() is called.
        The next() method of the returned iterator does not even need to worry about
        the termination condition (as StopIteration will be raised by DataSet.minibatches
        before an improper call to minibatches_nowrap's next() is made).
        That next() method can assert that its next row will always be within [0,len(dataset)).
        The iterator returned by minibatches_nowrap does not need to implement
        a next_index() method either, as this will be provided by MinibatchWrapAroundIterator.
        """
        raise AbstractFunction()

    def __len__(self):
        """
        len(dataset) returns the number of examples in the dataset.
        By default, a DataSet is a 'stream', i.e. it has an unbounded length (sys.maxint).
        Sub-classes which implement finite-length datasets should redefine this method.
        Some methods only make sense for finite-length datasets.
        """
        return sys.maxint

    def is_unbounded(self):
        """
        Tests whether a dataset is unbounded (e.g. a stream).
        """
        return len(self)==sys.maxint

    def hasFields(self,*fieldnames):
        """
        Return true if the given field name (or field names, if multiple arguments are
        given) is recognized by the DataSet (i.e. can be used as a field name in one
        of the iterators).

        The default implementation may be inefficient (O(# fields in dataset)), as it calls the fieldNames()
        method. Many datasets may store their field names in a dictionary, which would allow more efficiency.
        """
        return len(unique_elements_list_intersection(fieldnames,self.fieldNames()))>0
        
    def fieldNames(self):
        """
        Return the list of field names that are supported by the iterators,
        and for which hasFields(fieldname) would return True.
        """
        raise AbstractFunction()

    def __call__(self,*fieldnames):
        """
        Return a dataset that sees only the fields whose name are specified.
        """
        assert self.hasFields(*fieldnames)
        return self.fields(*fieldnames).examples()

    def fields(self,*fieldnames):
        """
        Return a DataSetFields object associated with this dataset.
        """
        return DataSetFields(self,*fieldnames)

    def __getitem__(self,i):
        """
        dataset[i] returns the (i+1)-th example of the dataset.
        dataset[i:j] returns the subdataset with examples i,i+1,...,j-1.
        dataset[i:j:s] returns the subdataset with examples i,i+2,i+4...,j-2.
        dataset[[i1,i2,..,in]] returns the subdataset with examples i1,i2,...,in.
        dataset['key'] returns a property associated with the given 'key' string.
        If 'key' is a fieldname, then the VStacked field values (iterable over
        field values) for that field is returned. Other keys may be supported
        by different dataset subclasses. The following key names are encouraged:
          - 'description': a textual description or name for the dataset
          - '<fieldname>.type': a type name or value for a given <fieldname>

        Note that some stream datasets may be unable to implement random access, i.e.
        arbitrary slicing/indexing
        because they can only iterate through examples one or a minibatch at a time
        and do not actually store or keep past (or future) examples.

        The default implementation of getitem uses the minibatches iterator
        to obtain one example, one slice, or a list of examples. It may not
        always be the most efficient way to obtain the result, especially if
        the data are actually stored in a memory array.
        """
        # check for an index
        if type(i) is int:
            return DataSet.MinibatchToSingleExampleIterator(
                self.minibatches(minibatch_size=1,n_batches=1,offset=i)).next()
        rows=None
        # or a slice
        if type(i) is slice:
            if not i.start: i.start=0
            if not i.step: i.step=1
            if i.step is 1:
                return self.minibatches(minibatch_size=i.stop-i.start,n_batches=1,offset=i.start).next().examples()
            rows = range(i.start,i.stop,i.step)
        # or a list of indices
        elif type(i) is list:
            rows = i
        if rows is not None:
            examples = [self[row] for row in rows]
            fields_values = zip(*examples)
            return MinibatchDataSet(
                Example(self.fieldNames(),[ self.valuesVStack(fieldname,field_values)
                                            for fieldname,field_values
                                            in zip(self.fieldNames(),fields_values)]),
                self.valuesVStack,self.valuesHStack)
        # else check for a fieldname
        if self.hasFields(i):
            return self.minibatches(fieldnames=[i],minibatch_size=len(self),n_batches=1,offset=0).next()[0]
        # else we are trying to access a property of the dataset
        assert i in self.__dict__ # else it means we are trying to access a non-existing property
        return self.__dict__[i]

    def valuesHStack(self,fieldnames,fieldvalues):
        """
        Return a value that corresponds to concatenating (horizontally) several field values.
        This can be useful to merge some fields. The implementation of this operation is likely
        to involve a copy of the original values. When the values are numpy arrays, the
        result should be numpy.hstack(values). If it makes sense, this operation should
        work as well when each value corresponds to multiple examples in a minibatch
        e.g. if each value is a Ni-vector and a minibatch of length L is a LxNi matrix,
        then the result should be a Lx(N1+N2+..) matrix equal to numpy.hstack(values).
        The default is to use numpy.hstack for numpy.ndarray values, and a list
        pointing to the original values for other data types.
        """
        all_numpy=True
        for value in fieldvalues:
            if not type(value) is numpy.ndarray:
                all_numpy=False
        if all_numpy:
            return numpy.hstack(fieldvalues)
        # the default implementation of horizontal stacking is to put values in a list
        return fieldvalues


    def valuesVStack(self,fieldname,values):
        """
        Return a value that corresponds to concatenating (vertically) several values of the
        same field. This can be important to build a minibatch out of individual examples. This
        is likely to involve a copy of the original values. When the values are numpy arrays, the
        result should be numpy.vstack(values).
        The default is to use numpy.vstack for numpy.ndarray values, and a list
        pointing to the original values for other data types.
        """
        all_numpy=True
        for value in values:
            if not type(value) is numpy.ndarray:
                all_numpy=False
        if all_numpy:
            return numpy.vstack(values)
        # the default implementation of vertical stacking is to put values in a list
        return values

    def __or__(self,other):
        """
        dataset1 | dataset2 returns a dataset whose list of fields is the concatenation of the list of
        fields of the argument datasets. This only works if they all have the same length.
        """
        return HStackedDataSet(self,other)

    def __and__(self,other):
        """
        dataset1 & dataset2 is a dataset that concatenates the examples from the argument datasets
        (and whose length is the sum of the length of the argument datasets). This only
        works if they all have the same fields.
        """
        return VStackedDataSet(self,other)

def hstack(datasets):
    """
    hstack(dataset1,dataset2,...) returns dataset1 | datataset2 | ...
    which is a dataset whose fields list is the concatenation of the fields
    of the individual datasets.
    """
    assert len(datasets)>0
    if len(datasets)==1:
        return datasets[0]
    return HStackedDataSet(datasets)

def vstack(datasets):
    """
    vstack(dataset1,dataset2,...) returns dataset1 & datataset2 & ...
    which is a dataset which iterates first over the examples of dataset1, then
    over those of dataset2, etc.
    """
    assert len(datasets)>0
    if len(datasets)==1:
        return datasets[0]
    return VStackedDataSet(datasets)

class FieldsSubsetDataSet(DataSet):
    """
    A sub-class of DataSet that selects a subset of the fields.
    """
    def __init__(self,src,fieldnames):
        self.src=src
        self.fieldnames=fieldnames
        assert src.hasFields(*fieldnames)
        self.valuesHStack = src.valuesHStack
        self.valuesVStack = src.valuesVStack

    def __len__(self): return len(self.src)
    
    def fieldNames(self):
        return self.fieldnames

    def __iter__(self):
        class FieldsSubsetIterator(object):
            def __init__(self,ds):
                self.ds=ds
                self.src_iter=ds.src.__iter__()
                self.example=None
            def __iter__(self): return self
            def next(self):
                complete_example = self.src_iter.next()
                if self.example:
                    self.example._values=[complete_example[field]
                                          for field in self.ds.fieldnames]
                else:
                    self.example=Example(self.ds.fieldnames,
                                         [complete_example[field] for field in self.ds.fieldnames])
                return self.example
        return FieldsSubsetIterator(self)

    def minibatches_nowrap(self,fieldnames,minibatch_size,n_batches,offset):
        assert self.hasFields(*fieldnames)
        return self.src.minibatches_nowrap(fieldnames,minibatch_size,n_batches,offset)
    def __getitem__(self,i):
        return FieldsSubsetDataSet(self.src[i],self.fieldnames)
    
        
class DataSetFields(LookupList):
    """
    Although a DataSet iterates over examples (like rows of a matrix), an associated
    DataSetFields iterates over fields (like columns of a matrix), and can be understood
    as a transpose of the associated dataset.

    To iterate over fields, one can do
    * for fields in dataset.fields()
    * for fields in dataset(field1,field2,...).fields() to select a subset of fields
    * for fields in dataset.fields(field1,field2,...) to select a subset of fields
    and each of these fields is iterable over the examples:
    * for field_examples in dataset.fields():
        for example_value in field_examples:
           ...
    but when the dataset is a stream (unbounded length), it is not recommanded to do 
    such things because the underlying dataset may refuse to access the different fields in
    an unsynchronized ways. Hence the fields() method is illegal for streams, by default.
    The result of fields() is a DataSetFields object, which iterates over fields,
    and whose elements are iterable over examples. A DataSetFields object can
    be turned back into a DataSet with its examples() method:
      dataset2 = dataset1.fields().examples()
    and dataset2 should behave exactly like dataset1 (in fact by default dataset2==dataset1).

    DataSetFields can be concatenated vertically or horizontally. To be consistent with
    the syntax used for DataSets, the | concatenates the fields and the & concatenates
    the examples.
    """
    def __init__(self,dataset,*fieldnames):
        original_dataset=dataset
        if not fieldnames:
            fieldnames=dataset.fieldNames()
        elif not fieldnames==dataset.fieldNames():
            dataset = FieldsSubsetDataSet(dataset,fieldnames)
        assert dataset.hasFields(*fieldnames)
        self.dataset=dataset

        if isinstance(dataset,MinibatchDataSet):
            LookupList.__init__(self,fieldnames,list(dataset._fields))
        elif isinstance(original_dataset,MinibatchDataSet):
            LookupList.__init__(self,fieldnames,
                                [original_dataset._fields[field]
                                 for field in fieldnames])
        else:
            minibatch_iterator = dataset.minibatches(fieldnames,
                                                     minibatch_size=len(dataset),
                                                     n_batches=1)
            minibatch=minibatch_iterator.next()
            LookupList.__init__(self,fieldnames,minibatch)
        
    def examples(self):
        return self.dataset
    
    def __or__(self,other):
        """
        fields1 | fields2 is a DataSetFields that whose list of examples is the concatenation
        of the list of examples of DataSetFields fields1 and fields2.
        """
        return (self.examples() + other.examples()).fields()

    def __and__(self,other):
        """
        fields1 + fields2 is a DataSetFields that whose list of fields is the concatenation
        of the fields of DataSetFields fields1 and fields2.
        """
        return (self.examples() | other.examples()).fields()

    
class MinibatchDataSet(DataSet):
    """
    Turn a LookupList of same-length fields into an example-iterable dataset.
    Each element of the lookup-list should be an iterable and sliceable, all of the same length.
    """
    def __init__(self,fields_lookuplist,values_vstack=DataSet().valuesVStack,
                 values_hstack=DataSet().valuesHStack):
        """
        The user can (and generally should) also provide values_vstack(fieldname,fieldvalues)
        and a values_hstack(fieldnames,fieldvalues) functions behaving with the same
        semantics as the DataSet methods of the same name (but without the self argument).
        """
        self._fields=fields_lookuplist
        assert len(fields_lookuplist)>0
        self.length=len(fields_lookuplist[0])
        for field in fields_lookuplist[1:]:
            assert self.length==len(field)
        self.values_vstack=values_vstack
        self.values_hstack=values_hstack

    def __len__(self):
        return self.length

    def __getitem__(self,i):
        if type(i) is int:
            return Example(self._fields.keys(),[field[i] for field in self._fields])
        if type(i) in (slice,list):
            return MinibatchDataSet(Example(self._fields.keys(),
                                            [field[i] for field in self._fields]),
                                    self.valuesVStack,self.valuesHStack)
        if self.hasFields(i):
            return self._fields[i]
        assert i in self.__dict__ # else it means we are trying to access a non-existing property
        return self.__dict__[i]

    def fieldNames(self):
        return self._fields.keys()

    def hasFields(self,*fieldnames):
        for fieldname in fieldnames:
            if fieldname not in self._fields.keys():
                return False
        return True

    def minibatches_nowrap(self,fieldnames,minibatch_size,n_batches,offset):
        class Iterator(object):
            def __init__(self,ds):
                self.ds=ds
                self.next_example=offset
                assert minibatch_size > 0
                if offset+minibatch_size > ds.length:
                    raise NotImplementedError()
            def __iter__(self):
                return self
            def next(self):
                upper = self.next_example+minibatch_size
                assert upper<=self.ds.length
                minibatch = Example(self.ds._fields.keys(),
                                    [field[self.next_example:upper]
                                     for field in self.ds._fields])
                self.next_example+=minibatch_size
                return minibatch

        return Iterator(self)

    def valuesVStack(self,fieldname,fieldvalues):
        return self.values_vstack(fieldname,fieldvalues)
    
    def valuesHStack(self,fieldnames,fieldvalues):
        return self.values_hstack(fieldnames,fieldvalues)
    
class HStackedDataSet(DataSet):
    """
    A DataSet that wraps several datasets and shows a view that includes all their fields,
    i.e. whose list of fields is the concatenation of their lists of fields.

    If a field name is found in more than one of the datasets, then either an error is
    raised or the fields are renamed (either by prefixing the __name__ attribute 
    of the dataset + ".", if it exists, or by suffixing the dataset index in the argument list).

    TODO: automatically detect a chain of stacked datasets due to A | B | C | D ...
    """
    def __init__(self,datasets,accept_nonunique_names=False,description=None,field_types=None):
        DataSet.__init__(self,description,field_types)
        self.datasets=datasets
        self.accept_nonunique_names=accept_nonunique_names
        self.fieldname2dataset={}

        def rename_field(fieldname,dataset,i):
            if hasattr(dataset,"__name__"):
                return dataset.__name__ + "." + fieldname
            return fieldname+"."+str(i)
            
        # make sure all datasets have the same length and unique field names
        self.length=None
        names_to_change=[]
        for i in xrange(len(datasets)):
            dataset = datasets[i]
            length=len(dataset)
            if self.length:
                assert self.length==length
            else:
                self.length=length
            for fieldname in dataset.fieldNames():
                if fieldname in self.fieldname2dataset: # name conflict!
                    if accept_nonunique_names:
                        fieldname=rename_field(fieldname,dataset,i)
                        names2change.append((fieldname,i))
                    else:
                        raise ValueError("Incompatible datasets: non-unique field name = "+fieldname)
                self.fieldname2dataset[fieldname]=i
        for fieldname,i in names_to_change:
            del self.fieldname2dataset[fieldname]
            self.fieldname2dataset[rename_field(fieldname,self.datasets[i],i)]=i
            
    def hasFields(self,*fieldnames):
        for fieldname in fieldnames:
            if not fieldname in self.fieldname2dataset:
                return False
        return True

    def fieldNames(self):
        return self.fieldname2dataset.keys()
            
    def minibatches_nowrap(self,fieldnames,minibatch_size,n_batches,offset):

        class HStackedIterator(object):
            def __init__(self,hsds,iterators):
                self.hsds=hsds
                self.iterators=iterators
            def __iter__(self):
                return self
            def next(self):
                # concatenate all the fields of the minibatches
                return reduce(LookupList.__add__,[iterator.next() for iterator in self.iterators])
                                     
        assert self.hasfields(fieldnames)
        # find out which underlying datasets are necessary to service the required fields
        # and construct corresponding minibatch iterators
        if fieldnames:
            datasets=set([])
            fields_in_dataset=dict([(dataset,[]) for dataset in datasets])
            for fieldname in fieldnames:
                dataset=self.datasets[self.fieldnames2dataset[fieldname]]
                datasets.add(dataset)
                fields_in_dataset[dataset].append(fieldname)
            datasets=list(datasets)
            iterators=[dataset.minibatches(fields_in_dataset[dataset],minibatch_size,n_batches,offset)
                       for dataset in datasets]
        else:
            datasets=self.datasets
            iterators=[dataset.minibatches(None,minibatch_size,n_batches,offset) for dataset in datasets]
        return HStackedIterator(self,iterators)


    def valuesVStack(self,fieldname,fieldvalues):
        return self.datasets[self.fieldname2dataset[fieldname]].valuesVStack(fieldname,fieldvalues)
    
    def valuesHStack(self,fieldnames,fieldvalues):
        """
        We will use the sub-dataset associated with the first fieldname in the fieldnames list
        to do the work, hoping that it can cope with the other values (i.e. won't care
        about the incompatible fieldnames). Hence this heuristic will always work if
        all the fieldnames are of the same sub-dataset.
        """
        return self.datasets[self.fieldname2dataset[fieldnames[0]]].valuesHStack(fieldnames,fieldvalues)

class VStackedDataSet(DataSet):
    """
    A DataSet that wraps several datasets and shows a view that includes all their examples,
    in the order provided. This clearly assumes that they all have the same field names
    and all (except possibly the last one) are of finite length.

    TODO: automatically detect a chain of stacked datasets due to A + B + C + D ...
    """
    def __init__(self,datasets):
        self.datasets=datasets
        self.length=0
        self.index2dataset={}
        assert len(datasets)>0
        fieldnames = datasets[-1].fieldNames()
        self.datasets_start_row=[]
        # We use this map from row index to dataset index for constant-time random access of examples,
        # to avoid having to search for the appropriate dataset each time and slice is asked for.
        for dataset,k in enumerate(datasets[0:-1]):
            assert dataset.is_unbounded() # All VStacked datasets (except possibly the last) must be bounded (have a length).
            L=len(dataset)
            for i in xrange(L):
                self.index2dataset[self.length+i]=k
            self.datasets_start_row.append(self.length)
            self.length+=L
            assert dataset.fieldNames()==fieldnames
        self.datasets_start_row.append(self.length)
        self.length+=len(datasets[-1])
        # If length is very large, we should use a more memory-efficient mechanism
        # that does not store all indices
        if self.length>1000000:
            # 1 million entries would require about 60 meg for the index2dataset map
            # TODO
            print "A more efficient mechanism for index2dataset should be implemented"

    def __len__(self):
        return self.length
    
    def fieldNames(self):
        return self.datasets[0].fieldNames()

    def hasFields(self,*fieldnames):
        return self.datasets[0].hasFields(*fieldnames)

    def locate_row(self,row):
        """Return (dataset_index, row_within_dataset) for global row number"""
        dataset_index = self.index2dataset[row]
        row_within_dataset = self.datasets_start_row[dataset_index]
        return dataset_index, row_within_dataset
        
    def minibatches_nowrap(self,fieldnames,minibatch_size,n_batches,offset):

        class VStackedIterator(object):
            def __init__(self,vsds):
                self.vsds=vsds
                self.next_row=offset
                self.next_dataset_index,self.next_dataset_row=self.vsds.locate_row(offset)
                self.current_iterator,self.n_left_at_the_end_of_ds,self.n_left_in_mb= \
                  self.next_iterator(vsds.datasets[0],offset,n_batches)

            def next_iterator(self,dataset,starting_offset,batches_left):
                L=len(dataset)
                ds_nbatches = (L-starting_offset)/minibatch_size
                if batches_left is not None:
                    ds_nbatches = max(batches_left,ds_nbatches)
                if minibatch_size>L:
                    ds_minibatch_size=L
                    n_left_in_mb=minibatch_size-L
                    ds_nbatches=1
                else:
                    n_left_in_mb=0
                return dataset.minibatches(fieldnames,minibatch_size,ds_nbatches,starting_offset), \
                       L-(starting_offset+ds_nbatches*minibatch_size), n_left_in_mb

            def move_to_next_dataset(self):
                if self.n_left_at_the_end_of_ds>0:
                    self.current_iterator,self.n_left_at_the_end_of_ds,self.n_left_in_mb= \
                      self.next_iterator(vsds.datasets[self.next_dataset_index],
                                         self.n_left_at_the_end_of_ds,1)
                else:
                    self.next_dataset_index +=1
                    if self.next_dataset_index==len(self.vsds.datasets):
                        self.next_dataset_index = 0
                    self.current_iterator,self.n_left_at_the_end_of_ds,self.n_left_in_mb= \
                      self.next_iterator(vsds.datasets[self.next_dataset_index],starting_offset,n_batches)
                
            def __iter__(self):
                return self

            def next(self):
                dataset=self.vsds.datasets[self.next_dataset_index]
                mb = self.next_iterator.next()
                if self.n_left_in_mb:
                    extra_mb = []
                    while self.n_left_in_mb>0:
                        self.move_to_next_dataset()
                        extra_mb.append(self.next_iterator.next())
                    mb = Example(fieldnames,
                                       [dataset.valuesVStack(name,
                                                             [mb[name]]+[b[name] for b in extra_mb])
                                            for name in fieldnames])
                    
                self.next_row+=minibatch_size
                self.next_dataset_row+=minibatch_size
                if self.next_row+minibatch_size>len(dataset):
                    self.move_to_next_dataset()
                return examples
        return VStackedIterator(self)
                        
class ArrayFieldsDataSet(DataSet):
    """
    Virtual super-class of datasets whose field values are numpy array,
    thus defining valuesHStack and valuesVStack for sub-classes.
    """
    def __init__(self,description=None,field_types=None):
        DataSet.__init__(self,description,field_types)
    def valuesHStack(self,fieldnames,fieldvalues):
        """Concatenate field values horizontally, e.g. two vectors
        become a longer vector, two matrices become a wider matrix, etc."""
        return numpy.hstack(fieldvalues)
    def valuesVStack(self,fieldname,values):
        """Concatenate field values vertically, e.g. two vectors
        become a two-row matrix, two matrices become a longer matrix, etc."""
        return numpy.vstack(values)

class ArrayDataSet(ArrayFieldsDataSet):
    """
    An ArrayDataSet stores the fields as groups of columns in a numpy tensor,
    whose first axis iterates over examples, second axis determines fields.
    If the underlying array is N-dimensional (has N axes), then the field
    values are (N-2)-dimensional objects (i.e. ordinary numbers if N=2).
    """

    def __init__(self, data_array, fields_columns):
        """
        Construct an ArrayDataSet from the underlying numpy array (data) and
        a map (fields_columns) from fieldnames to field columns. The columns of a field are specified
        using the standard arguments for indexing/slicing: integer for a column index,
        slice for an interval of columns (with possible stride), or iterable of column indices.
        """
        self.data=data_array
        self.fields_columns=fields_columns

        # check consistency and complete slices definitions
        for fieldname, fieldcolumns in self.fields_columns.items():
            if type(fieldcolumns) is int:
                assert fieldcolumns>=0 and fieldcolumns<data_array.shape[1]
            elif type(fieldcolumns) is slice:
                start,step=None,None
                if not fieldcolumns.start:
                    start=0
                if not fieldcolumns.step:
                    step=1
                if start or step:
                    self.fields_columns[fieldname]=slice(start,fieldcolumns.stop,step)
            elif hasattr(fieldcolumns,"__iter__"): # something like a list
                for i in fieldcolumns:
                    assert i>=0 and i<data_array.shape[1]

    def fieldNames(self):
        return self.fields_columns.keys()

    def __len__(self):
        return len(self.data)

    def __getitem__(self,i):
        """More efficient implementation than the default __getitem__"""
        fieldnames=self.fields_columns.keys()
        if type(i) is int:
            return Example(fieldnames,
                           [self.data[i,self.fields_columns[f]] for f in fieldnames])
        if type(i) in (slice,list):
            return MinibatchDataSet(Example(fieldnames,
                                            [self.data[i,self.fields_columns[f]] for f in fieldnames]),
                                    self.valuesVStack,self.valuesHStack)
        # else check for a fieldname
        if self.hasFields(i):
            return Example([i],[self.data[self.fields_columns[i],:]])
        # else we are trying to access a property of the dataset
        assert i in self.__dict__ # else it means we are trying to access a non-existing property
        return self.__dict__[i]
        
            
    def minibatches_nowrap(self,fieldnames,minibatch_size,n_batches,offset):
        class ArrayDataSetIterator(object):
            def __init__(self,dataset,fieldnames,minibatch_size,n_batches,offset):
                if fieldnames is None: fieldnames = dataset.fieldNames()
                # store the resulting minibatch in a lookup-list of values
                self.minibatch = LookupList(fieldnames,[0]*len(fieldnames))
                self.dataset=dataset
                self.minibatch_size=minibatch_size
                assert offset>=0 and offset<len(dataset.data)
                assert offset+minibatch_size<=len(dataset.data)
                self.current=offset
            def __iter__(self):
                return self
            def next(self):
                sub_data =  self.dataset.data[self.current:self.current+self.minibatch_size]
                self.minibatch._values = [sub_data[:,self.dataset.fields_columns[f]] for f in self.minibatch._names]
                self.current+=self.minibatch_size
                return self.minibatch

        return ArrayDataSetIterator(self,fieldnames,minibatch_size,n_batches,offset)


class CachedDataSet(DataSet):
  """
  Wrap a dataset whose values are computationally expensive to obtain
  (e.g. because they involve some computation, or disk access),
  so that repeated accesses to the same example are done cheaply,
  by caching every example value that has been accessed at least once.

  Optionally, for finite-length dataset, all the values can be computed
  (and cached) upon construction of the CachedDataSet, rather at the
  first access.

  TODO: add disk-buffering capability, so that when the cache becomes too
  big for memory, we cache things on disk, trying to keep in memory only
  the record most likely to be accessed next.
  """
  def __init__(self,source_dataset,cache_all_upon_construction=False):
      self.source_dataset=source_dataset
      self.cache_all_upon_construction=cache_all_upon_construction
      if cache_all_upon_construction:
          self.cached_examples = zip(*source_dataset.minibatches(minibatch_size=len(source_dataset)).__iter__().next())
      else:
          self.cached_examples = []

      self.fieldNames = source_dataset.fieldNames
      self.hasFields = source_dataset.hasFields
      self.valuesHStack = source_dataset.valuesHStack
      self.valuesVStack = source_dataset.valuesVStack
      
  def __len__(self):
      return len(self.source_dataset)

  def minibatches_nowrap(self,fieldnames,minibatch_size,n_batches,offset):
      class CacheIterator(object):
          def __init__(self,dataset):
              self.dataset=dataset
              self.current=offset
          def __iter__(self): return self
          def next(self):
              upper = self.current+minibatch_size
              cache_len = len(self.dataset.cached_examples)
              if upper>=cache_len: # whole minibatch is not already in cache
                  # cache everything from current length to upper
                  for example in self.dataset.source_dataset[cache_len:upper]:
                      self.dataset.cached_examples.append(example)
              all_fields_minibatch = Example(self.dataset.fieldNames(),
                                             self.dataset.cached_examples[self.current:self.current+minibatch_size])
              if self.dataset.fieldNames()==fieldnames:
                  return all_fields_minibatch
              return Example(fieldnames,[all_fields_minibatch[name] for name in fieldnames])
      return CacheIterator(self)

                      
class ApplyFunctionDataSet(DataSet):
  """
  A dataset that contains as fields the results of applying a given function
  example-wise or minibatch-wise to all the fields of an input dataset.
  The output of the function should be an iterable (e.g. a list or a LookupList)
  over the resulting values.

  In minibatch mode, the function is expected to work on minibatches (takes
  a minibatch in input and returns a minibatch in output). More precisely,
  it means that each element of the input or output list should be iterable
  and indexable over the individual example values (typically these
  elements will be numpy arrays). All of the elements in the input and
  output lists should have the same length, which is the length of the
  minibatch.

  The function is applied each time an example or a minibatch is accessed.
  To avoid re-doing computation, wrap this dataset inside a CachedDataSet.

  If the values_{h,v}stack functions are not provided, then
  the input_dataset.values{H,V}Stack functions are used by default.
  """
  def __init__(self,input_dataset,function,output_names,minibatch_mode=True,
               values_hstack=None,values_vstack,
               description=None,fieldtypes=None):
      """
      Constructor takes an input dataset that has as many fields as the function
      expects as inputs. The resulting dataset has as many fields as the function
      produces as outputs, and that should correspond to the number of output names
      (provided in a list).

      Note that the expected semantics of the function differs in minibatch mode
      (it takes minibatches of inputs and produces minibatches of outputs, as
      documented in the class comment).
      """
      self.input_dataset=input_dataset
      self.function=function
      self.output_names=output_names
      self.minibatch_mode=minibatch_mode
      DataSet.__init__(description,fieldtypes)
      self.valuesHStack = values_hstack if values_hstack else input_dataset.valuesHStack
      self.valuesVStack = values_vstack if values_vstack else input_dataset.valuesVStack

  def __len__(self):
      return len(self.input_dataset)

  def fieldnames(self):
      return self.output_names

  def minibatches_nowrap(self,fieldnames,minibatch_size,n_batches,offset):
      class ApplyFunctionIterator(object):
          def __init__(self,output_dataset):
              self.input_dataset=output_dataset.input_dataset
              self.output_dataset=output_dataset
              self.input_iterator=input_dataset.minibatches(minibatch_size=minibatch_size,
                                                            n_batches=n_batches,offset=offset).__iter__()

          def __iter__(self): return self

          def next(self):
              function_inputs = self.input_iterator.next()
              all_output_names = self.output_dataset.output_names
              if self.output_dataset.minibatch_mode:
                  function_outputs = self.output_dataset.function(function_inputs)
              else:
                  input_examples = zip(*function_inputs)
                  output_examples = [self.output_dataset.function(input_example)
                                     for input_example in input_examples]
                  function_outputs = [self.output_dataset.valuesVStack(name,values)
                                      for name,values in zip(all_output_names,
                                                             zip(*output_examples))]
              all_outputs = Example(all_output_names,function_outputs)
              if fieldnames==all_output_names:
                  return all_outputs
              return Example(fieldnames,[all_outputs[name] for name in fieldnames])

      return ApplyFunctionIterator(self.input_dataset,self)

  def __iter__(self): # only implemented for increased efficiency
      class ApplyFunctionSingleExampleIterator(object):
          def __init__(self,output_dataset):
              self.current=0
              self.output_dataset=output_dataset
              self.input_iterator=output_dataset.input_dataset.__iter__()
          def __iter__(self): return self
          def next(self):
              function_inputs = self.input_iterator.next()
              if self.output_dataset.minibatch_mode:
                  function_outputs = [output[0] for output in self.output_dataset.function(function_inputs)]
              else:
                  function_outputs = self.output_dataset.function(function_inputs)
              return Example(self.output_dataset.output_names,function_outputs)
      return ApplyFunctionSingleExampleIterator(self)
  

def supervised_learning_dataset(src_dataset,input_fields,target_fields,weight_field=None):
    """
    Wraps an arbitrary DataSet into one for supervised learning tasks by forcing the
    user to define a set of fields as the 'input' field and a set of fields
    as the 'target' field. Optionally, a single weight_field can also be defined.
    """
    args = ((input_fields,'input'),(output_fields,'target'))
    if weight_field: args+=(([weight_field],'weight'))
    return src_dataset.merge_fields(*args)