# 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_batchnormalisation, cnn_cifar_small, cnn_cifar_nodropout, \
    cnn_cifar_small_batchnormalisation
from models.DataLoader import DataLoader
from scripts.progress_bar import printProgress
from time import time, localtime
nb_epoch = 200
early_stopping = True
early_stopping_count = 0
early_stopping_wait = 3
train_loss = []
valid_loss = []
learning_rate = [0.0001, 0.001, 0.01]

# let's train the model using SGD + momentum (how original).
sgd = SGD(lr=learning_rate[-1], decay=1e-6, momentum=0.9, nesterov=True)
model.compile(loss='mean_squared_error', optimizer=sgd)

# load validation data from the h5py file (heavy lifting here)
x_valid, y_valid = loader.get_valid()
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, Y_batch = loader.next_train_batch()
        loss_list.append(model.train_on_batch(X_batch, Y_batch))

        # 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:', suffix='batch error: %0.5f, ETA: %0.2f sec.'%(np.array(loss_list).mean(), eta), barLength=50)

# 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.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_batchnormalisation, cnn_cifar_small, cnn_cifar_nodropout, \
    cnn_cifar_small_batchnormalisation
from models.DataLoader import DataLoader
from scripts.progress_bar import printProgress
from time import time, localtime

nb_epoch = 200
early_stopping = True
early_stopping_count = 0
early_stopping_wait = 3
train_loss = []
valid_loss = []
learning_rate = [0.0001, 0.001, 0.01]

# let's train the model using SGD + momentum (how original).
sgd = SGD(lr=learning_rate[-1], decay=1e-6, momentum=0.9, nesterov=True)
model.compile(loss='mean_squared_error', optimizer=sgd)

# load validation data from the h5py file (heavy lifting here)
x_valid, y_valid = loader.get_valid()
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, Y_batch = loader.next_train_batch()
        loss_list.append(model.train_on_batch(X_batch, Y_batch))

        # 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:', suffix='batch error: %0.5f, ETA: %0.2f sec.'%(np.array(loss_list).mean(), eta), barLength=50)

    printProgress(i, loader.n_iter_train - 1, prefix='Progress:', suffix='batch error: %0.5f' % (np.array(loss_list).mean()), barLength=50)