from fuel.datasets import Dataset
import numpy as np
import theano
floatX = theano.config.floatX

class Pylearn2Dataset(Dataset):
    '''Pylearn2Dataset wraps a `pylearn2.dataset` object and adds only the
    minimal `fuel` interface. An object of this class can be used as input to
    `fuel.streams.DataStream`.
    
    Parameters
    ----------
    dataset: `pylearn2.dataset` object 
        Note that this is expecting the actual the object will be initialized inside
    batch_size: int
        Batch size to be used by the `pylearn2.dataset` iterator.
    '''
    def __init__(self, dataset, batch_size, **kwargs):
        self.pylearn2_dataset = dataset
        self.sources = self.pylearn2_dataset.get_data_specs()[1]
        self.batch_size = batch_size 

    def open(self):
        num_examples = self.pylearn2_dataset.get_num_examples()
        iterator = self.pylearn2_dataset.iterator(
                   self.batch_size,
                   num_examples/self.batch_size,
                   mode='sequential',
                   data_specs=self.pylearn2_dataset.get_data_specs(),
                   return_tuple=True)
        return iterator    
    def get_data(self,state=None,request=None):
        return next(state)

class Pylearn2DatasetNoise(Dataset):
    '''Pylearn2DatasetNoise is the same as `Pylearn2Dataset` with some an 
    extra batch of random nubmer.
    
    Parameters
    ----------
    dataset: `pylearn2.dataset` object 
        Note that this is expecting the actual the object will be 
        initialized inside
    batch_size: int
        Batch size to be used by the `pylearn2.dataset` iterator.
    noise_dim: int
        Dimmension of the noise batch
    '''
    def __init__(self, dataset, batch_size, noise_dim, which_sources=[0,1], 
                **kwargs):
        self.pylearn2_dataset = dataset
        self.sources = self.pylearn2_dataset.get_data_specs()[1]
        self.sources = tuple([self.sources[i] for i in which_sources])
        self.sources = self.sources + tuple('eps')
        self.batch_size = batch_size 
        self.noise_dim = noise_dim
        self.which_sources = which_sources
    def open(self):
        num_examples = self.pylearn2_dataset.get_num_examples()
        iterator = self.pylearn2_dataset.iterator(
                   self.batch_size,
                   num_examples/self.batch_size,
                   mode='sequential',
                   data_specs=self.pylearn2_dataset.get_data_specs(),
                   return_tuple=True)
        return iterator    
    def get_data(self,state=None,request=None):
        batch = next(state)
        timelen = batch[0].shape[0]
        batch = tuple([batch[i] for i in self.which_sources])
        eps = np.random.normal(0,1,size=(timelen, 
                               self.batch_size, 
                               self.noise_dim)).astype(floatX)
        batch = batch + tuple(eps)
        return (batch,)

@@ -307,20 +307,16 @@
 tream%60.%0A
-    
 %0A    Par
@@ -368,33 +368,32 @@
 .dataset%60 object
- 
 %0A        Note th
@@ -599,16 +599,31 @@
 ch_size,
+ which_sources,
  **kwarg
@@ -748,32 +748,196 @@
 elf.
-batch_size = batch_size 
+sources = tuple(%5Bself.sources%5Bi%5D for i in which_sources%5D)%0A        self.sources = self.sources + tuple('eps')%0A        self.batch_size = batch_size%0A        self.which_sources = which_sources
 %0A%0A  
@@ -1315,36 +1315,32 @@
  return iterator
-    
 %0A    def get_dat
@@ -1380,22 +1380,23 @@
         
-return
+batch =
  next(st
@@ -1399,16 +1399,102 @@
 t(state)
+%0A        batch = tuple(%5Bbatch%5Bi%5D for i in self.which_sources%5D)%0A        return (batch,)
 %0A%0Aclass 
@@ -1597,17 +1597,16 @@
  some an
- 
 %0A    ext
@@ -1632,20 +1632,16 @@
 nubmer.%0A
-    
 %0A    Par
@@ -1701,17 +1701,16 @@
 %60 object
- 
 %0A       
@@ -1767,17 +1767,16 @@
  will be
- 
 %0A       
@@ -2021,17 +2021,16 @@
 s=%5B0,1%5D,
- 
 %0A       
@@ -2312,17 +2312,16 @@
 tch_size
- 
 %0A       
@@ -2780,20 +2780,16 @@
 iterator
-    
 %0A    def
@@ -3004,17 +3004,16 @@
 timelen,
- 
 %0A       
@@ -3052,17 +3052,16 @@
 ch_size,
- 
 %0A       

update pylearn2dataset