# Copyright 2015 Matthieu Courbariaux, Zhouhan Lin

"""
This file is adapted from BinaryConnect:

    https://github.com/MatthieuCourbariaux/BinaryConnect

Running this script should reproduce the results of a feed forward net trained
on MNIST.

To train a vanilla feed forward net with ordinary backprop:
   1. type "git checkout fullresolution" to switch to the "fullresolution" branch
   2. execute "python mnist.py"

To train a feed forward net with Binary Connect + quantized backprop:
   1. type "git checkout binary" to switch to the "binary" branch
   2. execute "python mnist.py"

To train a feed forward net with Ternary Connect + quantized backprop:
   1. type "git checkout ternary" to switch to the "ternary" branch
   2. execute "python mnist.py"

"""

import gzip
import cPickle
import numpy as np
import os
import os.path
import sys
import time

from trainer import Trainer
from model import Network
from layer import linear_layer, ReLU_layer  

from pylearn2.datasets.mnist import MNIST
from pylearn2.utils import serial
       

if __name__ == "__main__":
    
    rng = np.random.RandomState(1234)
    train_set_size = 50000
    
    # data augmentation
    zero_pad = 0
    affine_transform_a = 0
    affine_transform_b = 0
    horizontal_flip = False
    
    # batch
    # keep a multiple a factor of 10000 if possible
    # 10000 = (2*5)^4
    batch_size = 200
    number_of_batches_on_gpu = train_set_size/batch_size
    BN = True
    BN_epsilon=1e-4 # for numerical stability
    BN_fast_eval= True
    dropout_input = 1.
    dropout_hidden = 1.
    shuffle_examples = True
    shuffle_batches = False

    # Termination criteria
    n_epoch = 1000
    monitor_step = 2
    
    # LR 
    LR = .3
    LR_fin = .01
    LR_decay = (LR_fin/LR)**(1./n_epoch)    
    M= 0.
    
    # architecture
    n_inputs = 784
    n_units = 1024
    n_classes = 10
    n_hidden_layer = 3
    
    # BinaryConnect
    BinaryConnect = True
    stochastic = True
    
    # Old hyperparameters
    binary_training=False 
    stochastic_training=False
    binary_test=False
    stochastic_test=False
    if BinaryConnect == True:
        binary_training=True      
        if stochastic == True:   
            stochastic_training=True  
        else:
            binary_test=True
    
    print 'Loading the dataset' 
    
    train_set = MNIST(which_set= 'train', start=0, stop = train_set_size, center = True)
    valid_set = MNIST(which_set= 'train', start=50000, stop = 60000, center = True)
    test_set = MNIST(which_set= 'test', center = True)
    
    # bc01 format
    train_set.X = train_set.X.reshape(train_set_size,1,28,28)
    valid_set.X = valid_set.X.reshape(10000,1,28,28)
    test_set.X = test_set.X.reshape(10000,1,28,28)
    
    # flatten targets
    train_set.y = np.hstack(train_set.y)
    valid_set.y = np.hstack(valid_set.y)
    test_set.y = np.hstack(test_set.y)
    
    # Onehot the targets
    train_set.y = np.float32(np.eye(10)[train_set.y])    
    valid_set.y = np.float32(np.eye(10)[valid_set.y])
    test_set.y = np.float32(np.eye(10)[test_set.y])
    
    # for hinge loss
    train_set.y = 2* train_set.y - 1.
    valid_set.y = 2* valid_set.y - 1.
    test_set.y = 2* test_set.y - 1.
    
    print 'Creating the model'
    
    class PI_MNIST_model(Network):

        def __init__(self, rng):
            
            Network.__init__(self, n_hidden_layer = n_hidden_layer, BN = BN)
            
            print "    Fully connected layer 1:"
            self.layer.append(ReLU_layer(rng = rng, n_inputs = n_inputs, n_units = n_units,
                BN = BN, BN_epsilon=BN_epsilon, dropout=dropout_input,
                binary_training=binary_training, stochastic_training=stochastic_training,
                binary_test=binary_test, stochastic_test=stochastic_test))
            
            for k in range(n_hidden_layer-1):
                
                print "    Fully connected layer "+ str(k) +":"
                self.layer.append(ReLU_layer(rng = rng, n_inputs = n_units, n_units = n_units,
                    BN = BN, BN_epsilon=BN_epsilon, dropout=dropout_hidden, 
                    binary_training=binary_training, stochastic_training=stochastic_training,
                    binary_test=binary_test, stochastic_test=stochastic_test))
                
            print "    L2 SVM layer:"
            self.layer.append(linear_layer(rng = rng, n_inputs = n_units, n_units = n_classes,
                BN = BN, BN_epsilon=BN_epsilon, dropout=dropout_hidden, 
                binary_training=binary_training, stochastic_training=stochastic_training,
                binary_test=binary_test, stochastic_test=stochastic_test))
    
    model = PI_MNIST_model(rng = rng)
    
    print 'Creating the trainer'
    
    trainer = Trainer(rng = rng,
        train_set = train_set, valid_set = valid_set, test_set = test_set,
        model = model, load_path = None, save_path = None,
        zero_pad=zero_pad,
        affine_transform_a=affine_transform_a, # a is (more or less) the rotations
        affine_transform_b=affine_transform_b, # b is the translations
        horizontal_flip=horizontal_flip,
        LR = LR, LR_decay = LR_decay, LR_fin = LR_fin,
        M = M,
        BN = BN, BN_fast_eval=BN_fast_eval,
        batch_size = batch_size, number_of_batches_on_gpu = number_of_batches_on_gpu,
        n_epoch = n_epoch, monitor_step = monitor_step,
        shuffle_batches = shuffle_batches, shuffle_examples = shuffle_examples)
    
    print 'Building'
    
    trainer.build()
    
    print 'Training'
    
    start_time = time.clock()  
    trainer.train()
    end_time = time.clock()
    print 'The training took %i seconds'%(end_time - start_time)
    
    print 'Display weights'
    
    import matplotlib.pyplot as plt
    import matplotlib.cm as cm
    from filter_plot import tile_raster_images
    
    W = np.transpose(model.layer[0].W.get_value())
    W = tile_raster_images(W,(28,28),(4,4),(2, 2))
    plt.imshow(W, cmap = cm.Greys_r)
    plt.savefig(core_path + '_features.png')
