import pylearn.algorithms.stacker as models_stacker
import pylearn.algorithms.regressor as models_reg
import theano
import numpy
import time

class StackBinRegressor(models_reg.BinRegressor):
    def __init__(self, input = None, target = None, regularize = True):
        super(StackBinRegressor, self).__init__(input, target, regularize)
        self.build_extensions()

def test_train(mode = theano.Mode('c|py', 'fast_run')):
    reg = models_stacker.Stacker([(StackBinRegressor, 'output'),
                                  (StackBinRegressor, 'output')],
                                  regularize = False)
    #print reg.global_update[1].pretty(mode = mode.excluding('inplace'))

    model = reg.make([100, 200, 1],
                     lr = 0.01,
                     mode = mode,
                     seed = 10)

    R = numpy.random.RandomState(100)
    t1 = time.time()
    for i in xrange(1001):
        data = R.random_integers(0, 1, size = (10, 100))
        targets = data[:, 6].reshape((10, 1))
        cost = model.update(data, targets)
        if i % 100 == 0:
            print i, '\t', cost, '\t', 1*(targets.T == model.classify(data).T)
    t2 = time.time()
    return t2 - t1

if __name__ == '__main__':
    print 'optimized:'
    t1 = test_train(theano.Mode('c|py', 'fast_run'))
    print 'time:',t1
    print

    print 'not optimized:'
    t2 = test_train(theano.Mode('c|py', 'fast_compile'))
    print 'time:',t2