def run_epoch(trainobj):
    """
    Runs an epoch. Returns True to continue or
    False to terminate.
    """

    if trainobj.first_callbacks_and_monitoring:
        trainobj.run_callbacks_and_monitoring()
        trainobj.first_callbacks_and_monitoring = False
        return True

    rval = True
    if trainobj.algorithm is None:
        rval = trainobj.model.train_all(dataset=trainobj.dataset)
        if rval is not None:
            raise ValueError("Model.train_all should not return " +
                             "anything. Use Model.continue_learning " +
                             "to control whether learning continues.")
        rval = post_epoch(trainobj)
    else:
        with log_timing(logger, None, level=logging.DEBUG,
                        callbacks=[trainobj.total_seconds.set_value]):
            with log_timing(logger, None, final_msg='Time this epoch:',
                            callbacks=[trainobj.training_seconds.set_value]):
                rval = trainobj.algorithm.train(dataset=trainobj.dataset)
            if rval is not None:
                raise ValueError("TrainingAlgorithm.train should not "
                                 "return anything. Use "
                                 "TrainingAlgorithm.continue_learning "
                                 "to control whether learning "
                                 "continues.")
            rval = post_epoch(trainobj)
    return rval

def process_dataset(model, dataset, data_specs=None, output_fn=None, batch_size=128):
    
    if data_specs is None:
        data_specs = (CompositeSpace((
                                model.get_input_space(), 
                                model.get_output_space())), 
                           ("features", "targets"));
    
    if output_fn is None:                
        with log_timing(log, 'compiling output_fn'):         
            minibatch = model.get_input_space().make_theano_batch();
            output_fn = theano.function(inputs=[minibatch], 
                                        outputs=model.fprop(minibatch));
    
    it = dataset.iterator(mode='sequential',
                          batch_size=batch_size,
                          data_specs=data_specs);
    y_pred = [];
    y_real = [];                
    output = [];
    for minibatch, target in it:
        out = output_fn(minibatch); # this hangs for convnet on Jeep2
        output.append(out);
        # print out
        # print out.shape
        y_pred.append(np.argmax(out, axis = 1));
        y_real.append(np.argmax(target, axis = 1));
    y_pred = np.hstack(y_pred);
    y_real = np.hstack(y_real);  
    output = np.vstack(output);
    
    return y_real, y_pred, output;

 def on_monitor(self, model, dataset, algorithm):
     
     epoch = algorithm.monitor._epochs_seen;
     model_file = self.save_path + self.save_prefix + str(epoch) + '.pkl'; 
     
     with log_timing(log, 'saving model to {}'.format(model_file)):
         serial.save(model_file, model, on_overwrite = 'backup')

    def on_monitor(self, model, dataset, algorithm):
        """
        Looks whether the model performs better than earlier
        - or equally good (modification).
        If it's the case, saves the model.

        Parameters
        ----------
        model : pylearn2.models.model.Model
            model.monitor must contain a channel with name given by
            self.channel_name
        dataset : pylearn2.datasets.dataset.Dataset
            Not used
        algorithm : TrainingAlgorithm
            Not used
        """
        monitor = model.monitor
        channels = monitor.channels
        channel = channels[self.channel_name]
        val_record = channel.val_record
        new_cost = val_record[-1]

        if self.coeff * new_cost <= self.coeff * self.best_cost and \
           monitor._epochs_seen >= self.start_epoch:
            self.best_cost = new_cost
            # Update the tag of the model object before saving it.
            self._update_tag(model)
            if self.store_best_model:
                self.best_model = deepcopy(model)
            if self.save_path is not None:
                with log_timing(log, 'Saving to ' + self.save_path):
                    serial.save(self.save_path, model, on_overwrite='backup')

def load_results(experiment_root):
    # load the model (mlp_best.pkl)
    model_file = os.path.join(experiment_root, 'mlp_best.pkl');    
    with log_timing(log, 'loading model from {}'.format(model_file)):  
        model = serial.load(model_file);    

    # load train
    train_yaml_file = os.path.join(experiment_root, 'train.yaml');
    train_yaml = load_yaml_template(train_yaml_file);
    
    # fix dataset path
    localizer = PathLocalizer();
    train_yaml = localizer.localize_yaml(train_yaml);
    
    with log_timing(log, 'loading train from {}'.format(train_yaml_file)):      
        train = load_yaml(train_yaml)[0];
    
    return train, model;

def extract_output(experiment_root):
    train, model = load_results(experiment_root);
        
    # get the datasets with their names from the monitor
    for key, dataset in train.algorithm.monitoring_dataset.items():
        # process each dataset 
        with log_timing(log, 'processing dataset \'{}\''.format(key)): 
            y_real, y_pred, output = process_dataset(model, dataset)
            
            save(os.path.join(experiment_root, 'cache', key+'_output.pklz'), (y_real, y_pred, output));    

def train_mlp(params):
    train, yaml_str = load_yaml_file(
                   os.path.join(os.path.dirname(__file__), 'cross_trial_template.yaml'),
                   params=params,
                   );
    
    save_yaml_file(yaml_str, os.path.join(params.experiment_root, 'settings.yaml'));
        
    with log_timing(log, 'training network'):    
        train.main_loop();

def save_yaml_file(yaml_str, yaml_file_path):
    if save_yaml_file is not None:
        with log_timing(log, 'saving yaml to {}'.format(yaml_file_path)):
            save_dir = os.path.dirname(yaml_file_path);
            if save_dir == '':
                save_dir = '.'
            if not os.path.exists(save_dir):
                os.makedirs(save_dir)
            with  open(yaml_file_path, 'w') as yaml_file:
                yaml_file.write(yaml_str) 
            yaml_file.close();

def train_convnet(config):
    
    train, yaml_str = load_yaml_file(
                   os.path.join(os.path.dirname(__file__), 'train_convnet_template.yaml'),
                   params=config,
                   );
    
    save_yaml_file(yaml_str, os.path.join(config.experiment_root, 'settings.yaml'));
        
    with log_timing(log, 'training network'):    
        train.main_loop();

 def __init__(self, filepath):
     self.filepath = filepath
     with log_timing(log, 'loading data from {}'.format(filepath)):
         tmp = load(filepath)
         if len(tmp) == 2:
             self.data, self.metadata = tmp
             self.targets = None
         elif len(tmp) == 3:
             self.data, self.metadata, self.targets = tmp
         else:
             raise ValueError('got {} objects instead of 2 or 3.'.format(len(tmp)))

def load_data_file(filename):

    #data = np.loadtxt(filename, dtype=float, delimiter=' ', skiprows=1); #, autostrip=True, names=False) 
    with log_timing(log, 'loading data from {}'.format(filename)):
        data = np.genfromtxt(filename,  dtype=theano.config.floatX, delimiter=' ', skip_header=1, autostrip=True);    
    log.info('loaded {}'.format(data.shape));
    
#     print data.shape;
#     print data[0];
#     print data[-1];

    return data;

def load_yaml(yaml_template, params=None):    
    print params;
    
    if params is not None:
        yaml_str = yaml_template % params;
    else:
        yaml_str = yaml_template;
    print yaml_str;

    with log_timing(log, 'parsing yaml'):    
        obj = yaml_parse.load(yaml_str);
    
    return obj, yaml_str;

def load_yaml(yaml_template, params=None):    
    log.debug('params: {}'.format(params))
    
    if params is not None:
        yaml_str = yaml_template % params
    else:
        yaml_str = yaml_template
    log.debug(yaml_str)

    with log_timing(log, 'parsing yaml'):    
        obj = yaml_parse.load(yaml_str)
    
    return obj, yaml_str

    def main_loop(self):
        """
        Repeatedly runs an epoch of the training algorithm, runs any
        epoch-level callbacks, and saves the model.
        """
        if self.algorithm is None:
            self.model.monitor = Monitor.get_monitor(self.model)
            self.setup_extensions()
            self.run_callbacks_and_monitoring()
            while True:
                rval = self.model.train_all(dataset=self.dataset)
                if rval is not None:
                    raise ValueError("Model.train_all should not return anything. Use Model.continue_learning to control whether learning continues.")
                self.model.monitor.report_epoch()
                if self.save_freq > 0 and self.model.monitor.epochs_seen % self.save_freq == 0:
                    self.save()
                continue_learning = self.model.continue_learning()
                assert continue_learning in [True, False, 0, 1]
                if not continue_learning:
                    break
        else:
            self.algorithm.setup(model=self.model, dataset=self.dataset)
            self.setup_extensions()
            if not hasattr(self.model, 'monitor'):
                # TODO: is this really necessary? I just put this error here
                # to prevent an AttributeError later, but I think we could
                # rewrite to avoid the AttributeError
                raise RuntimeError("The algorithm is responsible for setting"
                        " up the Monitor, but failed to.")
            self.run_callbacks_and_monitoring()
            while True:
                with log_timing(log, None, final_msg='Time this epoch:'):
                    rval = self.algorithm.train(dataset=self.dataset)
                if rval is not None:
                    raise ValueError("TrainingAlgorithm.train should not return anything. Use TrainingAlgorithm.continue_learning to control whether learning continues.")
                self.model.monitor.report_epoch()
                self.run_callbacks_and_monitoring()
                if self.save_freq > 0 and self.model.monitor._epochs_seen % self.save_freq == 0:
                    self.save()
                continue_learning =  self.algorithm.continue_learning(self.model)
                assert continue_learning in [True, False, 0, 1]
                if not continue_learning:
                    break

        self.model.monitor.training_succeeded = True

        if self.save_freq > 0:
            self.save()

    def __init__(self, save_dir):
        PYLEARN2_TRAIN_DIR = preprocess('${PYLEARN2_TRAIN_DIR}')
        PYLEARN2_TRAIN_BASE_NAME = preprocess('${PYLEARN2_TRAIN_BASE_NAME}')

        src = os.path.join(PYLEARN2_TRAIN_DIR, PYLEARN2_TRAIN_BASE_NAME)
        dst = os.path.join(save_dir, PYLEARN2_TRAIN_BASE_NAME)

        if not os.path.exists(save_dir):
            os.makedirs(save_dir)
        if os.path.exists(save_dir) and not os.path.isdir(save_dir):
            raise IOError("save path %s exists, not a directory" % save_dir)
        elif not os.access(save_dir, os.W_OK):
            raise IOError("permission error creating %s" % dst)

        with log_timing(log, 'copying yaml from {} to {}'.format(src, dst)):
            copyfile(src, dst)

 def run_one_epoch(self, datasets, remember_best):
     batch_generator = self.iterator.get_batches(datasets['train'],
         shuffle=True)
     with log_timing(log, None, final_msg='Time updates following epoch:'):
         for inputs, targets in batch_generator:
             if self.batch_modifier is not None:
                 inputs, targets = self.batch_modifier.process(inputs,
                     targets)
             # could happen that batch modifier has removed all inputs...
             if len(inputs) > 0:
                 self.train_func(inputs, targets)
     
     self.monitor_epoch(datasets)
     self.print_epoch()
     if remember_best:
         self.remember_extension.remember_epoch(self.monitor_chans,
             self.all_params)

 def save(self):
     """Saves the model."""
     # TODO-- save state of training algorithm so training can be
     # resumed after a crash
     for extension in self.extensions:
         extension.on_save(self.model, self.dataset, self.algorithm)
     if self.save_path is not None:
         with log_timing(log, "Saving to " + self.save_path):
             if self.first_save and (not self.allow_overwrite) and os.path.exists(self.save_path):
                 # Every job overwrites its own output on the second save
                 # and every save thereafter. The "allow_overwrite" flag
                 # only pertains to overwriting the output of previous jobs.
                 raise IOError("Trying to overwrite file when not allowed.")
             try:
                 # Make sure that saving does not serialize the dataset
                 self.dataset._serialization_guard = SerializationGuard()
                 serial.save(self.save_path, self.model, on_overwrite="backup")
             finally:
                 self.dataset._serialization_guard = None
         self.first_save = False

def train_mlp(params):
    
#     sda_file = os.path.join(params.experiment_root, 'sda', 'sda_all.pkl');

    # check whether pre-trained SDA is there
    pretrained = True;
    for i in xrange(len(params.hidden_layers_sizes)):
        sda_layer_file = params.get(('layer{}_content').format(i));
        if not os.path.isfile(sda_layer_file):
            log.info('did not find pre-trained SDA layer model at {}. re-computing SDA'.format(sda_layer_file));
            pretrained = False;
            break;
        else:
            log.info('found pre-trained SDA layer model at {}'.format(sda_layer_file));
    
    if not pretrained:
        train_sda(params);
        
    n_layers = len(params.hidden_layers_sizes);
        
    if params.learning_rule == 'AdaDelta':
        yaml_template = 'train_sda_mlp_template.AdaDelta.yaml'
    else:
        if n_layers == 3:
            yaml_template = 'train_sda_mlp_template.Momentum.yaml'
        elif n_layers == 2:
            yaml_template = 'train_sda_mlp_template.Momentum.2layers.yaml'
        else:
            raise '{} layers not supported'.format(n_layers);
    
    train, train_yaml_str = load_yaml_file(
                   os.path.join(os.path.dirname(__file__), yaml_template),
                   params=params,
                   );
                   
    save_yaml_file(train_yaml_str, os.path.join(params.experiment_root, 'mlp_train.yaml'));
    
    with log_timing(log, 'training MLP'):    
        train.main_loop();
        
    log.info('done');

def extract_output(config, best_epoch):
    # load best model
    model_file = os.path.join(config.experiment_root, "epochs", "epoch{}.pkl".format(best_epoch))
    print "loading " + model_file
    model = serial.load(model_file)

    #     print model;

    # additional dataset params
    config.start_sample = 11200
    config.stop_sample = 12800
    config.name = "test"

    # load dataset
    dataset, dataset_yaml = load_yaml_file(
        os.path.join(os.path.dirname(__file__), "..", "run", "dataset_template.yaml"), params=config
    )

    with log_timing(log, "processing dataset"):
        y_real, y_pred, output = process_dataset(model, dataset)

    return y_real, y_pred, output

def split_trial(path, trial_len):
    
    log.info('processing {}'.format(path));
    
    datafile = glob.glob(os.path.join(path,'*.txt'))[0];
    metafile = glob.glob(os.path.join(path,'*_Trials_Onsets.xlsx'))[0];
    
    log.debug('data file: {}'.format(datafile));
    log.debug('meta file: {}'.format(metafile));

    onsets = load_xlsx_meta_file(metafile);    
    data = load_data_file(datafile);
    log.debug(onsets);
    
    onsets.append([len(data), 'end']); # artificial last marker

    trials = {};
    for i in xrange(len(onsets) - 1):
        onset, label = onsets[i];
        next_onset = onsets[i+1][0];
        
        # rounding to integers
        onset = int(math.floor(float(onset)));
        next_onset = int(math.floor(float(next_onset)));
        
        next_onset = min(onset+trial_len, next_onset);
        
        log.debug('[{}..{}) -> {}'.format(onset, next_onset, label));
        trial_data = np.vstack(data[onset:next_onset]);
        log.debug('{} samples extracted'.format(trial_data.shape));
        
        trials[label] = trial_data;
        
    filename = os.path.join(path, 'trials.pklz');
    with log_timing(log, 'saving to {}'.format(filename)):
        save(filename, trials);
        
    return trials;