class Cifar10(SingleBlock):

    def __init__(self, flatten=False, **kwargs):

        im_dir = os.environ['MOZI_DATA_PATH'] + '/cifar10/'
        path = 'http://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz'
        im_dir = get_file(fpath="{}/cifar-10-python.tar.gz".format(im_dir), origin=path, untar=True)
        self.label_names = ['airplane', 'automobile', 'bird', 'cat', 'deer',
                            'dog', 'frog','horse','ship','truck']

        self.img_shape = (3,32,32)
        self.img_size = np.prod(self.img_shape)
        self.n_classes = 10
        fnames = ['data_batch_%i' % i for i in range(1,6)] + ['test_batch']

        X = []
        y = []
        for fname in fnames:
            data_path = "{}/{}".format(im_dir, fname)
            with open(data_path) as fin:
                data_batch = cPickle.load(fin)
                if flatten:
                    X.extend(data_batch['data'].reshape((len(data_batch['data']), self.img_size)))
                else:
                    X.extend(data_batch['data'].reshape((len(data_batch['data']),)+self.img_shape))
                y.extend(data_batch['labels'])


        X_npy = np.array(X, dtype=floatX)
        X_npy /= 255.0
        y_npy = make_one_hot(y, onehot_size=self.n_classes)
        ridx = np.arange(len(X_npy))
        np.random.shuffle(ridx)
        X_npy = X_npy[ridx]
        y_npy = y_npy[ridx]

        super(Cifar10, self).__init__(X=X_npy, y=y_npy, **kwargs) class Sequential(Model):

    def __init__(self, input_var, output_var, verbose=True):
        """
        PARAM:
            input_var (T.vector() | T.matrix() | T.tensor3() | T.tensor4()):
                    The tensor variable input to the model that corresponds to
                    the number of dimensions of the input X of dataset
            input_var (T.vector() | T.matrix() | T.tensor3() | T.tensor4()):
                    The tensor variable output from the model that corresponds to
                    the number of dimensions of the output y of dataset
            verbose (bool):
                    print out the layer stats from each layer if True

        """
        self.input_var = input_var
        self.output_var = output_var
        self.layers = []
        self.verbose = verbose

    def add(self, layer):
        self.layers.append(layer)

    def pop(self, index):
        return self.layers.pop(index)

    def test_fprop(self, input_state, layers=None):
        test_layers_stats = []
        if layers is None:
            layers = xrange(len(self.layers))
        for i in layers:
            layer_output = self.layers[i]._test_fprop(input_state)
            stats = []
            if self.verbose:
                stats = self.layers[i]._layer_stats(input_state, layer_output)
            input_state = layer_output
            class_name = self.layers[i].__class__.__name__
            stats = [(str(i)+'_'+class_name+'_'+a, b) for (a,b) in stats]
            test_layers_stats += stats

        return input_state, test_layers_stats


    def train_fprop(self, input_state, layers=None):
        train_layers_stats = []
        if layers is None:
            layers = xrange(len(self.layers))
        for i in layers:
            layer_output = self.layers[i]._train_fprop(input_state)
            stats = []
            if self.verbose:
                stats = self.layers[i]._layer_stats(input_state, layer_output)
            input_state = layer_output
            class_name = self.layers[i].__class__.__name__
            stats = [(str(i)+'_'+class_name+'_'+a, b) for (a,b) in stats]
            train_layers_stats += stats

        return input_state, train_layers_stats


    def fprop(self, input_values):
        return self.fprop_layers(input_values)


    def fprop_layers(self, input_values, layers=None):
        output, stats = self.test_fprop(self.input_var, layers)
        if isinstance(self.input_var, (list, tuple)):
            f = theano.function(self.input_var, output, on_unused_input='warn', allow_input_downcast=True)
        else:
            f = theano.function([self.input_var], output, on_unused_input='warn', allow_input_downcast=True)

        if isinstance(input_values, tuple):
            return f(*input_values)
        else:
            return f(input_values)


    def get_layers(self):
        return self.layers class Dropout(Template):

    def __init__(self, dropout_below=0.5):
        '''
        PARAMS:
            dropout_below(float): probability of the inputs from the layer below been masked out
        '''
        self.dropout_below = dropout_below
        self.params = []


    def _test_fprop(self, state_below):
        """
        DESCRIPTION:
            resize the weight during testing for models trained with dropout.
            The weight will be resized to W' = self.dropout_below * W
        """
        return state_below * (1 - self.dropout_below)


    def _train_fprop(self, state_below):
        """
        DESCRIPTION:
            Applies dropout to the layer during training
        """
        return theano_rand.binomial(size=state_below.shape, n=1,
                                    p=(1-self.dropout_below),
                                    dtype=floatX) * state_below class Flatten(Template):

    def _train_fprop(self, state_below):
        size = T.prod(state_below.shape) / state_below.shape[0]
        nshape = (state_below.shape[0], size)
        return T.reshape(state_below, nshape) class Log:

    def __init__(self, experiment_name="experiment", description=None,
                save_outputs=False, save_model=False,
                save_epoch_error=False, save_to_database=None, logger=None):

        self.experiment_name = experiment_name
        self.description = description
        self.save_outputs = save_outputs
        self.save_model = save_model
        self.save_epoch_error = save_epoch_error
        self.save_to_database = save_to_database

        dt = datetime.now()
        dt = dt.strftime('%Y%m%d_%H%M_%S%f')

        self.exp_id = experiment_name + '_' + dt

        if save_outputs or save_model:
            save_dir = os.environ['MOZI_SAVE_PATH']
            if not os.path.exists(save_dir):
                os.mkdir(save_dir)

            self.exp_dir = save_dir + '/' + self.exp_id
            if not os.path.exists(self.exp_dir):
                os.mkdir(self.exp_dir)

        self.logger = logger
        if self.logger is None:
            self.logger = logging.getLogger(__name__)
            self.logger.setLevel(logging.DEBUG)

        self.logger.info('exp_id: ' + experiment_name)

        if save_outputs:
            ch = logging.FileHandler(filename=self.exp_dir+'/outputs.log')
            ch.setLevel(logging.INFO)
            formatter = logging.Formatter('%(message)s')
            ch.setFormatter(formatter)
            self.logger.addHandler(ch)

        if save_epoch_error:
            self.epoch_error_path = self.exp_dir+'/epoch_error.csv'
            with open(self.epoch_error_path, 'wb') as epoch_file:
                epoch_file.write('Epoch,Train_Cost,Valid_Cost,Valid_Error\n')

        if description is not None:
            self.logger.info('Description: ' + self.description)

        if save_to_database:
            self.first_time_record = True
            if not os.path.exists(os.environ['MOZI_DATABASE_PATH']):
                os.mkdir(os.environ['MOZI_DATABASE_PATH'])

    def info(self, msg):
        self.logger.info(msg)

    def print_records(self):
        sorted_ls = sorted(self.save_to_database['records'].iteritems(),
                             key=operator.itemgetter(0))
        for key, value in sorted_ls:
            self.info(key + ': ' + str(value))

    def _log_outputs(self, outputs):
        dt = datetime.now()
        dt = dt.strftime('%Y-%m-%d %H:%M')
        self.logger.info('Time: ' + dt)

        for (name, val) in outputs:
            self.logger.info(name + ': ' + str(val))

        if self.save_outputs:
            self.logger.info('[ outputs saved to: %s ]\n' %self.exp_id)

    def _save_model(self, model):
        with open(self.exp_dir+'/model.pkl', 'wb') as pkl_file:
            cPickle.dump(model, pkl_file)

    def _save_epoch_error(self, epoch, train_cost, valid_cost, valid_error):
        with open(self.epoch_error_path, 'ab') as epoch_file:
            epoch_file.write('{},{},{},{}\n'.format(epoch, train_cost, valid_cost, valid_error))

    def _save_to_database(self, epoch, train_cost, valid_error, best_valid_error):
        conn = sqlite3.connect(os.environ['MOZI_DATABASE_PATH'] + '/' + self.save_to_database['name'])
        cur = conn.cursor()

        if self.first_time_record:
            query = 'CREATE TABLE IF NOT EXISTS ' + self.experiment_name + \
                    '(exp_id TEXT PRIMARY KEY NOT NULL,'

            for k,v in self.save_to_database['records'].items():
                if type(v) is str:
                    query += k + ' TEXT,'
                elif type(v) is int:
                    query += k + ' INT,'
                elif type(v) is float:
                    query += k + ' REAL,'
                else:
                    try:
                        self.save_to_database['records'][k] = str(v)
                        query += str(k) + ' TEXT,'
                    except:
                        raise Exception("Error: The input types for records '{}' of {}".format(k, type(v))
                                    + " is not primitive types (str, int, float) and not castable as str.")

            query += 'epoch INT, train_cost REAL, valid_error REAL, best_valid_error REAL);'

            cur.execute(query)

            try:
                query = 'INSERT INTO ' + self.experiment_name + ' VALUES('
                ls = [self.exp_id]
                for k, v in self.save_to_database['records'].items():
                    query += '?,'
                    ls.append(v)
                query += '?,?,?,?,?);'
                ls.extend([epoch, train_cost, valid_error, best_valid_error])
                cur.execute(query, ls)
                self.first_time_record = False

            except sqlite3.OperationalError as err:
                self.logger.error('sqlite3.OperationalError: ' + err.message)
                self.logger.error('Solution: Change the experiment_name in Log() to a new name, '
                        + 'or drop the same table name from the database. '
                        + 'experiment_name is used as the table name.')
                raise

        else:
            cur.execute('UPDATE ' + self.experiment_name + ' SET ' +
                        'epoch = ?, ' +
                        'train_cost = ?,' +
                        'valid_error = ?,' +
                        'best_valid_error = ?' +
                        "WHERE exp_id='%s'"%self.exp_id,
                        [epoch,
                        train_cost,
                        valid_error,
                        best_valid_error])
        conn.commit()
        conn.close() class TrainObject():

    def __init__(self, model, dataset, train_cost, valid_cost, learning_method, stop_criteria, log=None, verbose=True):
        self.model = model
        self.dataset = dataset
        self.train_cost = train_cost
        self.valid_cost = valid_cost
        self.learning_method = learning_method
        self.stop_criteria = stop_criteria
        self.log = log
        self.verbose = verbose


        if self.log is None:
            # use default Log setting
            self.log = Log(logger=internal_logger)

        elif self.log.save_to_database:
            self.log.print_records()
            self.log.info('\n')


    def setup(self):

        self.log.info( '..begin setting up train object')

        #===================[ build params and deltas list ]==================#

        params = []
        deltas = []

        for i, layer in enumerate(self.model.layers):
            layer_name = "{}_{}".format(layer.__class__.__name__, i)
            if hasattr(layer, 'params'):
                for param in layer.params:
                    # checked that the param to be updated is shared variable
                    if is_shared_var(param):
                        param.name = str(i) + '_' + str(param.name)
                        param.name += '_' + layer.__class__.__name__
                        params += [param]
                        deltas += [shared_zeros(shape=param.shape.eval())]

        #=====================[ training params updates ]=====================#

        self.log.info("..update params: " + str(params))
        train_y_pred, train_layers_stats = self.model.train_fprop(self.model.input_var)
        train_cost = self.train_cost(self.model.output_var, train_y_pred).astype(floatX)
        gparams = T.grad(train_cost, params)
        train_updates = self.learning_method.update(deltas, params, gparams)

        #=================[ append updates from each layer ]==================#

        for i, layer in enumerate(self.model.layers):
            layer_name = "{}_{}".format(layer.__class__.__name__, i)
            if hasattr(layer, 'updates') and len(layer.updates) > 0:
                self.log.info("..{}: has shared variable updates".format(layer_name))
                train_updates += layer.updates

        #----[ append updates of stats from each layer to train updates ]-----#

        self.train_stats_names, train_stats_vars = split_list(train_layers_stats)
        train_stats_vars = [var.astype(floatX) for var in train_stats_vars]
        self.train_stats_shared = generate_shared_list(train_stats_vars)
        train_stats_updates = merge_lists(self.train_stats_shared, train_stats_vars)
        if self.verbose:
            train_updates += train_stats_updates

        #-------------------------[ train functions ]-------------------------#

        self.log.info('..begin compiling functions')
        self.training = theano.function(inputs=merge_var(self.model.input_var, self.model.output_var),
                                        outputs=train_cost,
                                        updates=train_updates,
                                        on_unused_input='warn',
                                        allow_input_downcast=True)

        self.log.info('..training function compiled')

        #=============================[ testing ]=============================#

        test_y_pred, test_layers_stats = self.model.test_fprop(self.model.input_var)

        #-----[ append updates of stats from each layer to test updates ]-----#

        self.test_stats_names, test_stats_vars = split_list(test_layers_stats)
        test_stats_vars = [var.astype(floatX) for var in test_stats_vars]
        self.test_stats_shared = generate_shared_list(test_stats_vars)
        test_stats_updates = []
        if self.verbose:
            test_stats_updates = merge_lists(self.test_stats_shared, test_stats_vars)

        #-------------------------[ test functions ]--------------------------#

        test_stopping_error = self.valid_cost(self.model.output_var, test_y_pred).astype(floatX)
        test_cost = self.train_cost(self.model.output_var, test_y_pred).astype(floatX)

        self.testing = theano.function(inputs=merge_var(self.model.input_var, self.model.output_var),
                                       outputs=(test_stopping_error, test_cost),
                                       updates=test_stats_updates,
                                       on_unused_input='warn',
                                       allow_input_downcast=True)

        self.log.info('..testing function compiled')


    def run(self):

        best_valid_error = float(sys.maxint)
        valid_error = float(sys.maxint)

        train_cost = float(sys.maxint)
        valid_cost = float(sys.maxint)

        train_stats_values = []
        valid_stats_values = []

        epoch = 0
        error_dcr = 0
        self.best_epoch_last_update = 0
        self.best_valid_last_update = float(sys.maxint)

        train_stats_names = ['train_' + name for name in self.train_stats_names]
        valid_stats_names = ['valid_' + name for name in self.test_stats_names]

        job_start = time.time()

        while (self.continue_learning(epoch, error_dcr, best_valid_error)):

            if epoch > 0:
                self.log.info("best_epoch_last_update: %d"%self.best_epoch_last_update)
                self.log.info("valid_error_decrease: %f"%error_dcr)
                self.log.info("best_valid_last_update: %f"%self.best_valid_last_update)
                self.log.info("========[ End of Epoch ]========\n\n")

            epoch += 1

            start_time = time.time()

            num_train_examples = 0
            total_train_cost = 0.
            train_stats_values = np.zeros(len(train_stats_names), dtype=floatX)

            num_valid_examples = 0
            total_valid_cost = 0.
            total_valid_stopping_cost = 0.
            valid_stats_values = np.zeros(len(valid_stats_names), dtype=floatX)

            blk = 0

            for block in self.dataset:
                block_time = time.time()
                blk += 1

                train_set = block.get_train()
                valid_set = block.get_valid()

                #====================[ Training Progress ]====================#
                if train_set.dataset_size > 0:
                    self.log.info('..training '+ self.dataset.__class__.__name__
                                + ' block %s/%s'%(blk, self.dataset.nblocks))

                    progbar = Progbar(target=train_set.dataset_size)
                    blk_sz = 0
                    for idx in train_set:
                        cost = self.training(*train_set[idx])
                        total_train_cost += cost * len(idx)
                        num_train_examples += len(idx)
                        train_stats_values += len(idx) * get_shared_values(self.train_stats_shared)
                        blk_sz += len(idx)
                        progbar.update(blk_sz)
                    print

                #===================[ Validating Progress ]===================#
                if valid_set.dataset_size > 0:

                    self.log.info('..validating ' + self.dataset.__class__.__name__
                                + ' block %s/%s'%(blk, self.dataset.nblocks))

                    progbar = Progbar(target=valid_set.dataset_size)
                    blk_sz = 0
                    for idx in valid_set:
                        stopping_cost, cost = self.testing(*valid_set[idx])
                        total_valid_cost += cost * len(idx)
                        total_valid_stopping_cost += stopping_cost * len(idx)
                        num_valid_examples += len(idx)
                        valid_stats_values += len(idx) * get_shared_values(self.test_stats_shared)
                        blk_sz += len(idx)
                        progbar.update(blk_sz)
                    print

                self.log.info('block time: %0.2fs'%(time.time()-block_time))
                self.log.info(get_mem_usage())

            #-------[ Update train best cost and error values ]-------#
            if num_train_examples > 0:
                train_cost = total_train_cost / num_train_examples
                train_stats_values /= num_train_examples

            #-------[ Update valid best cost and error values ]-------#
            if num_valid_examples > 0:
                valid_error = total_valid_stopping_cost / num_valid_examples
                valid_cost = total_valid_cost / num_valid_examples
                valid_stats_values /= num_valid_examples

                if valid_error < best_valid_error:
                    best_valid_error = valid_error
                    self.log.info('..best validation error so far')
                    if self.log.save_model:
                        self.log._save_model(self.model)
                        self.log.info('..model saved')

                if valid_error < self.best_valid_last_update:
                    error_dcr = self.best_valid_last_update - valid_error
                else:
                    error_dcr = 0

            #==============[ save to database, save epoch error]==============#
            if self.log.save_to_database:
                self.log._save_to_database(epoch, train_cost, valid_error, best_valid_error)
                self.log.info('..sent to database: %s:%s' % (self.log.save_to_database['name'],
                                                             self.log.experiment_name))

            if self.log.save_epoch_error:
                self.log._save_epoch_error(epoch, train_cost, valid_cost, valid_error)
                self.log.info('..epoch error saved')

            end_time = time.time()

            #=====================[ log outputs to file ]=====================#
            merged_train = merge_lists(train_stats_names, train_stats_values)
            merged_valid = merge_lists(valid_stats_names, valid_stats_values)

            outputs = [('epoch', epoch),
                        ('runtime(s)', int(end_time-start_time)),
                        ('train_' + self.train_cost.func_name, train_cost),
                        ('valid_' + self.train_cost.func_name, valid_cost),
                        ('valid_' + self.valid_cost.func_name, valid_error),
                        ('best_valid_' + self.valid_cost.func_name, best_valid_error)]

            outputs += merged_train + merged_valid
            self.log._log_outputs(outputs)


        job_end = time.time()
        self.log.info('Job Completed on %s'%time.strftime("%a, %d %b %Y %H:%M:%S", time.gmtime(job_end)))
        ttl_time = int(job_end - job_start)
        dt = datetime.timedelta(seconds=ttl_time)
        self.log.info('Total Time Taken: %s'%str(dt))
        self.log.info("========[ End of Job ]========\n\n")


    def continue_learning(self, epoch, error_dcr, best_valid_error):

        if epoch > self.stop_criteria['max_epoch']:
            return False

        elif self.stop_criteria['percent_decrease'] is None or \
            self.stop_criteria['epoch_look_back'] is None:
            return True

        elif np.abs(float(error_dcr) / self.best_valid_last_update) \
            >= self.stop_criteria['percent_decrease']:
            self.best_valid_last_update = best_valid_error
            self.best_epoch_last_update = epoch
            return True

        elif epoch - self.best_epoch_last_update > \
            self.stop_criteria['epoch_look_back']:
            return False

        else:
            return True def error(y, y_pred):
    L = T.neq(y_pred.argmax(axis=1), y.argmax(axis=1))
    return T.mean(L) def entropy(y, y_pred):
    y_pred = T.clip(y_pred, epsilon, 1.0 - epsilon)
    L = -(y * T.log(y_pred) + (1-y) * T.log(1-y_pred))
    return T.mean(L)
    # L = - T.sum(y * T.log(y_pred) + (1-y) * T.log(1-y_pred), axis=1)
    # return T.mean(L) def setenv():
    NNdir = os.path.dirname(os.path.realpath(__file__))
    NNdir = os.path.dirname(NNdir)

    # directory to save all the dataset
    if not os.getenv('MOZI_DATA_PATH'):
        os.environ['MOZI_DATA_PATH'] = NNdir + '/data'

    # directory for saving the database that is used for logging the results
    if not os.getenv('MOZI_DATABASE_PATH'):
        os.environ['MOZI_DATABASE_PATH'] = NNdir + '/database'

    # directory to save all the trained models and outputs
    if not os.getenv('MOZI_SAVE_PATH'):
        os.environ['MOZI_SAVE_PATH'] = NNdir + '/save'

    print('MOZI_DATA_PATH = ' + os.environ['MOZI_DATA_PATH'])
    print('MOZI_SAVE_PATH = ' + os.environ['MOZI_SAVE_PATH'])
    print('MOZI_DATABASE_PATH = ' + os.environ['MOZI_DATABASE_PATH']) def valid(x, y, kernel, stride):
    return ((x-kernel)/stride + 1, (y-kernel)/stride + 1) def full(x, y, kernel, stride):
    return ((x+kernel)/stride - 1, (y+kernel)/stride - 1)
import numpy as np
import os
import theano.tensor as T
from mozi.datasets.cifar10 import Cifar10
from mozi.model import Sequential
from mozi.layers.linear import *
from mozi.layers.noise import Dropout
from mozi.layers.activation import *
from mozi.layers.convolution import *
from mozi.layers.misc import Flatten
from mozi.layers.normalization import *
from mozi.log import Log
from mozi.train_object import TrainObject
from mozi.cost import error, entropy
from mozi.learning_method import *
from mozi.env import setenv
from mozi.utils.cnn_utils import valid, full



def train():
    batch_size = 256
    short_memory = 0.9
    learning_rate = 0.005
    data = Cifar10(batch_size=batch_size, train_valid_test_ratio=[4,1,1])
    _, c, h, w = data.train.X.shape

    model = Sequential(input_var=T.tensor4(), output_var=T.matrix())
    model.add(Convolution2D(input_channels=c, filters=8, kernel_size=(3,3), stride=(1,1), border_mode='full'))
    h, w = full(h, w, kernel=3, stride=1)
    model.add(BatchNormalization(dim=8, layer_type='conv', short_memory=short_memory))
    model.add(RELU())
    model.add(Convolution2D(input_channels=8, filters=16, kernel_size=(3,3), stride=(1,1), border_mode='valid'))
    h, w = valid(h, w, kernel=3, stride=1)
    model.add(BatchNormalization(dim=16, layer_type='conv', short_memory=short_memory))
    model.add(RELU())
    model.add(Pooling2D(poolsize=(4, 4), stride=(4,4), mode='max'))
    h, w = valid(h, w, kernel=4, stride=4)
    model.add(Flatten())
    model.add(Linear(16*h*w, 512))
    model.add(BatchNormalization(dim=512, layer_type='fc', short_memory=short_memory))
    model.add(RELU())

    model.add(Linear(512, 10))
    model.add(Softmax())

    # learning_method = RMSprop(learning_rate=learning_rate)
    learning_method = Adam(learning_rate=learning_rate)
    # learning_method = SGD(learning_rate=0.001)

    # Build Logger
    log = Log(experiment_name = 'cifar10_cnn_tutorial',
              description = 'This is a tutorial',
              save_outputs = True, # log all the outputs from the screen
              save_model = True, # save the best model
              save_epoch_error = True, # log error at every epoch
              save_to_database = {'name': 'hyperparam.sqlite3',
                                  'records': {'Batch_Size': batch_size,
                                              'Learning_Rate': learning_method.learning_rate}}
             ) # end log

    # put everything into the train object
    train_object = TrainObject(model = model,
                               log = log,
                               dataset = data,
                               train_cost = entropy,
                               valid_cost = error,
                               learning_method = learning_method,
                               stop_criteria = {'max_epoch' : 100,
                                                'epoch_look_back' : 10,
                                                'percent_decrease' : 0.01}
                               )
    # finally run the code
    train_object.setup()
    train_object.run()

    # test the model on test set
    ypred = model.fprop(data.get_test().X)
    ypred = np.argmax(ypred, axis=1)
    y = np.argmax(data.get_test().y, axis=1)
    accuracy = np.equal(ypred, y).astype('f4').sum() / len(y)
    print 'test accuracy:', accuracy


if __name__ == '__main__':

    setenv()