# UCF Senior Design 2017-18
# Group 38

from sklearn.model_selection import train_test_split as tts
import pickle
import tflearn
import os
import yaml

DATASET = 'leavens.pkl'
OUTPUT_FOLDER = '/output/buff_bobo'
SEED = 3

# Residual blocks
# 32 layers: n=5, 56 layers: n=9, 110 layers: n=18
n = 5

# Real-time data preprocessing
img_prep = tflearn.ImagePreprocessing()
img_prep.add_featurewise_zero_center(per_channel=True)

# Real-time data augmentation
img_aug = tflearn.ImageAugmentation()
img_aug.add_random_flip_leftright()

with open(DATASET, 'rb') as f:
    pickle = pickle.load(f, encoding='latin1')

x = pickle['x']
y = pickle['y']
(X, X_test, Y, Y_test) = tts(x, y, test_size=0.3, random_state=SEED)

X, Y = tflearn.data_utils.shuffle(X, Y)

net = tflearn.input_data(shape=[None, 448, 448, 3],
                         data_preprocessing=img_prep,
                         data_augmentation=img_aug)

net = tflearn.conv_2d(net, 64, 3, regularizer='L2', weight_decay=0.0001)
net = tflearn.residual_block(net, n, 64)
net = tflearn.residual_block(net, 1, 128, downsample=True)
net = tflearn.residual_block(net, n-1, 128)
net = tflearn.residual_block(net, 1, 256, downsample=True)
net = tflearn.residual_block(net, n-1, 256)
net = tflearn.residual_block(net, 1, 512, downsample=True)
net = tflearn.residual_block(net, n-1, 512)
net = tflearn.batch_normalization(net)
net = tflearn.activation(net, 'relu')
net = tflearn.global_avg_pool(net)

net = tflearn.fully_connected(net, 2, activation='softmax')
mom = tflearn.Momentum(0.1, lr_decay=0.1, decay_step=32000, staircase=True)
net = tflearn.regression(net, optimizer='adam',
                         loss='categorical_crossentropy')

# Training
model = tflearn.DNN(net, checkpoint_path=OUTPUT_FOLDER,
                    max_checkpoints=10, tensorboard_verbose=3,
                    clip_gradients=0.)

model.fit(X, Y, n_epoch=200, validation_set=(X_test, Y_test),
          show_metric=True, batch_size=32, shuffle=True,
          run_id='buff_bobo')

model_name = "buff_bobo"  + ".tfl"
model_path = os.path.join(OUTPUT_FOLDER, model_name)
model.save(model_path)