# Path: scripts/progress_bar.py
# def printProgress (iteration, total, prefix = '', suffix = '', decimals = 1, barLength = 100):
#     """
#     Call in a loop to create terminal progress bar
#     @params:
#         iteration   - Required  : current iteration (Int)
#         total       - Required  : total iterations (Int)
#         prefix      - Optional  : prefix string (Str)
#         suffix      - Optional  : suffix string (Str)
#         decimals    - Optional  : positive number of decimals in percent complete (Int)
#         barLength   - Optional  : character length of bar (Int)
#     """
#     formatStr = "{0:." + str(decimals) + "f}"
#     percent = formatStr.format(100 * (iteration / float(total)))
#     filledLength = int(round(barLength * iteration / float(total)))
#     bar = '#' * filledLength + '-' * (barLength - filledLength)
#     sys.stdout.write('\r%s |%s| %s%s %s' % (prefix, bar, percent, '%', suffix)),
#     if iteration == total:
#         sys.stdout.write('\n')
#     sys.stdout.flush()


import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from keras.optimizers import SGD
from models.CNN_CIFAR import cnn_cifar_small_batchnormalisation_class_scale
from models.DataLoader import DataLoader
from keras.models import load_model
from scripts.progress_bar import printProgress
from time import time, localtime

# model.compile(loss='categorical_crossentropy', optimizer='rmsprop',
#               loss_weights=[1., 1., 1.])

# load validation data from the h5py file (heavy lifting here)
x_valid, translation_x_class_valid, translation_y_class_valid,  y_scale_class_valid =\
    loader.get_valid_class()
best_valid = np.inf

for e in range(nb_epoch):
    print("epoch %d" % e)
    loss_list = []
    time_list = []
    time_start = time()
    for i in range(loader.n_iter_train):
        time_start_batch = time()
        X_batch, translation_x_class, translation_y_class, y_scale_class = \
            loader.next_train_batch_with_scale_buckets()
        l1, l2, l3, l4 = model.train_on_batch({'main_input': X_batch},
                                              {'translation_x_class': translation_x_class,
                                               'translation_y_class': translation_y_class,
                                               'y_scale_class': y_scale_class})
        loss_list.append(l1)
        translation_x_loss.append(l2)
        translation_y_loss.append(l3)
        scale_loss.append(l4)

        # calculate some time information
        time_list.append(time() - time_start_batch)
        eta = (loader.n_iter_train - i) * np.array(time_list).mean()

# Path: scripts/progress_bar.py
# def printProgress (iteration, total, prefix = '', suffix = '', decimals = 1, barLength = 100):
#     """
#     Call in a loop to create terminal progress bar
#     @params:
#         iteration   - Required  : current iteration (Int)
#         total       - Required  : total iterations (Int)
#         prefix      - Optional  : prefix string (Str)
#         suffix      - Optional  : suffix string (Str)
#         decimals    - Optional  : positive number of decimals in percent complete (Int)
#         barLength   - Optional  : character length of bar (Int)
#     """
#     formatStr = "{0:." + str(decimals) + "f}"
#     percent = formatStr.format(100 * (iteration / float(total)))
#     filledLength = int(round(barLength * iteration / float(total)))
#     bar = '#' * filledLength + '-' * (barLength - filledLength)
#     sys.stdout.write('\r%s |%s| %s%s %s' % (prefix, bar, percent, '%', suffix)),
#     if iteration == total:
#         sys.stdout.write('\n')
#     sys.stdout.flush()


# Path: CNN_training_with_scale.py
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from keras.optimizers import SGD
from models.CNN_CIFAR import cnn_cifar_small_batchnormalisation_class_scale
from models.DataLoader import DataLoader
from keras.models import load_model
from scripts.progress_bar import printProgress
from time import time, localtime


# model.compile(loss='categorical_crossentropy', optimizer='rmsprop',
#               loss_weights=[1., 1., 1.])

# load validation data from the h5py file (heavy lifting here)
x_valid, translation_x_class_valid, translation_y_class_valid,  y_scale_class_valid =\
    loader.get_valid_class()
best_valid = np.inf

for e in range(nb_epoch):
    print("epoch %d" % e)
    loss_list = []
    time_list = []
    time_start = time()
    for i in range(loader.n_iter_train):
        time_start_batch = time()
        X_batch, translation_x_class, translation_y_class, y_scale_class = \
            loader.next_train_batch_with_scale_buckets()
        l1, l2, l3, l4 = model.train_on_batch({'main_input': X_batch},
                                              {'translation_x_class': translation_x_class,
                                               'translation_y_class': translation_y_class,
                                               'y_scale_class': y_scale_class})
        loss_list.append(l1)
        translation_x_loss.append(l2)
        translation_y_loss.append(l3)
        scale_loss.append(l4)

        # calculate some time information
        time_list.append(time() - time_start_batch)
        eta = (loader.n_iter_train - i) * np.array(time_list).mean()

        printProgress(i, loader.n_iter_train-1, prefix='Progress:',