#!/usr/bin/env python
# The goal of this script is to load VGG layers, freeze them and use these build a autoencoder based on that

import os
import time
import cv2
import sys
import numpy as np
import glob
from utils import display_results_RGB, enumerate2
from keras.callbacks import EarlyStopping, ModelCheckpoint, TensorBoard
from network_utils import create_model_wider_filter_network, compute_weights
# from keras.preprocessing.image import ImageDataGenerator
from keras.preprocessing.image import load_img, img_to_array, list_pictures
from keras import optimizers
from keras.models import load_model

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"   # see issue #152
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

# This scripts tries to replicate the textures present in the training images
starting_dir = os.getcwd()

# master_folder = '/home/gcx/lstm_sequences/autoencoder-200/'
master_folder = '/home/gcx/lstm_sequences/anomaly_forth/'

os.chdir(starting_dir)

def image_generator(list_of_files_x, list_of_files_y, img_size, to_grayscale=False):
    small_img_size = 34
    while True:
        filename_idx = np.random.randint(0, len(list_of_files_x))
        filename_x = list_of_files_x[filename_idx]
        filename_y = list_of_files_y[filename_idx]

        try:
            x = img_to_array(load_img(filename_x, to_grayscale)) / 255.

            y = img_to_array(load_img(filename_y))[:, :, 0] / 255.
            x = cv2.resize(x,(img_size, img_size), interpolation=cv2.INTER_NEAREST)
            # y = cv2.resize(y,(img_size, img_size), interpolation=cv2.INTER_NEAREST)
            # _weights = compute_weights(y.reshape(1, img_size * img_size), img_size)
            # _weights = _weights.reshape((img_size*img_size))
            # # y = np.reshape(y, (y.shape[0], y.shape[1], 1))
            # y = y.reshape((img_size*img_size, 1, 1))

            y = cv2.resize(y, (small_img_size, small_img_size), interpolation=cv2.INTER_NEAREST)
            _weights = compute_weights(y.reshape(1, small_img_size * small_img_size), small_img_size)
            _weights = _weights.reshape((small_img_size* small_img_size))
            # y = np.reshape(y, (y.shape[0], y.shape[1], 1))
            y = y.reshape((small_img_size * small_img_size, 1, 1))
        except:
            print('Error on loading routine!')
            return

        yield (x, y, _weights)


def group_by_batch(dataset, batch_size):
    while True:
        try:
            sources, targets, weights = zip(*[next(dataset) for i in xrange(batch_size)])
            batch = (np.stack(sources), np.stack(targets), np.stack(weights))
            yield batch
        except:
            print('error on grouping routine')
            return


def load_dataset(directory, batch_size, _img_size, stage):
    print('directory + stage + x_:', directory + stage + '/x_' + stage)
    print('directory + stage + y_:', directory + stage + '/y_' + stage)
    x_files = list_pictures(directory + stage + '/x_' + stage)
    y_files = list_pictures(directory + stage + '/y_' + stage)

    generator = image_generator(x_files, y_files, _img_size)
    generator = group_by_batch(generator, batch_size)
    return generator


img_size = 540

# try to load a model
# model_to_load = glob.glob('models/vgg/rgb_540_reformulated/*.h5')
# if len(model_to_load) > 1:
#     sys.exit('Multiple files in folder. Do not know what to choose. Exiting!')
#
# elif len(model_to_load) == 1:
#     print('Found one file. going to load: ', model_to_load[0])
#     new_model = load_model(model_to_load[0])
#
# else:
#     print('No model found! going to define a new model')


new_model = create_model_wider_filter_network(img_size, 34)
# # optimizer = optimizers.adadelta(lr=0.1)
new_model.compile(optimizer='adadelta', loss='binary_crossentropy', sample_weight_mode='temporal', )
# new_model.compile(optimizer='adadelta', loss='mean_absolute_error', sample_weight_mode='temporal', )


current_date = time.strftime('%Y%m%d')
current_time = time.strftime('%H%M')
model_fn = 'models/ae_vgg_nt_based_dg' + '.ep{epoch:02d}-ls{val_loss:.5f}' + '_' \
           + current_date + '_' + current_time + '_0.h5'

cb_modelCheckpoint = ModelCheckpoint(model_fn, monitor='val_loss',
                     verbose=0, save_best_only=True,
                     save_weights_only=False, mode='min', period=1)

cb_EarlyStopping = EarlyStopping(monitor='val_loss', min_delta=0, patience=10, verbose=1, mode='min')

# TRAIN
batch_size = 10

nb_epoch = 30

dataset_dir = '/home/gcx/lstm_sequences/anomaly_forth/'

file_list = list_pictures(dataset_dir+'train/x_train/')
test_file_list = list_pictures(dataset_dir+'test/x_test/')

train_set = load_dataset(dataset_dir, batch_size, img_size, stage='train')
test_set = load_dataset(dataset_dir, batch_size, img_size, stage='test')
file_list_number = len(file_list)
test_file_list_number = len(test_file_list)

nb_steps_per_epoch = file_list_number/float(batch_size)
nb_validation_steps = test_file_list_number / float(batch_size)

new_model.fit_generator(train_set,
                        validation_data=test_set,
                        steps_per_epoch=nb_steps_per_epoch,
                        nb_epoch=nb_epoch,
                        callbacks=[TensorBoard(log_dir='/tmp/wider_0'),
                                   cb_modelCheckpoint],
                        validation_steps=nb_validation_steps, )

print('finished training')

# # LOAD DATA
# print('loading test data..')
# x_test, y_test = load_data(master_folder, 'test/', img_size, '*.tiff')
# print('loading terminated.')
#
# os.chdir(starting_dir)
#
# # normalize images
# x_test, _ = normalize_data(x_test, img_size)
#
# # normalize ground truth
# y_test, y_test_vectorized = normalize_data(y_test, img_size)
#
# # replicate images along channels dimension
# x_test_rgb = np.zeros((x_test.shape[0],x_test.shape[1],x_test.shape[2],3))
# x_test_rgb = np.stack((x_test[:, :, :, 0], x_test[:, :, :, 0], x_test[:, :, :, 0]), axis=3)
#
# # create sample weights for Y_TEST
# sample_weights_test = compute_weights(y_test_vectorized, img_size)
#
# # Evalute and get prediction
# print('result of evaluation: ', new_model.evaluate(x_test_rgb, y_test_vectorized, batch_size, 1))
# reconstructed_imgs = new_model.predict(x_test_rgb)
#
# n = 100  # how many digits we will display
# step_size = 23
#
# figure_name = 'fig1_vgg_ae_dg.png'
# display_results(x_test, y_test, reconstructed_imgs, img_size, step_size, n, figure_name)
#
# cv2.namedWindow('test')
# cv2.waitKey(0)
# print('finished script!')


tests_number = 100  # how many digits we will display
step_size = 15
channel_number = 3

# Load image one by one
os.chdir(master_folder + 'test/' + 'x_' + 'test')
x = np.zeros((tests_number, img_size, img_size, channel_number))

y = np.ones((tests_number, img_size, img_size))

# print(sorted(glob.glob('*.png'), key=lambda name: int(name[:-4])))
ptr = 0
for i, file_ptr in enumerate2(glob.glob('*.tiff'), start=0, step=step_size, stop= step_size * tests_number):
    _img_tmp = cv2.imread(file_ptr)
    _img_tmp = cv2.resize(_img_tmp, (img_size, img_size), interpolation=cv2.INTER_NEAREST) / 255.
    # _img_tmp = np.average(np.asarray(_img_tmp, np.float), axis=2)
    # print('max', np.max(np.max(_img_tmp)))
    x[ptr, :, :, :] = _img_tmp

    # print(_folder + stage + 'y_' + stage + os.path.split(file_ptr)[1])
    _gt_tmp = cv2.imread(master_folder + 'test/' + 'y_' + 'test/' + os.path.split(file_ptr)[1])
    _gt_tmp = cv2.resize(_gt_tmp, (img_size, img_size), interpolation=cv2.INTER_NEAREST)
    _gt_tmp = np.average(np.asarray(_gt_tmp, np.float), axis=2) / 255.
    # print('gt max', np.max(np.max(_gt_tmp)))
    y[ptr, :, :] = _gt_tmp

    ptr += 1

print('i, ptr:', i, ptr)

os.chdir(starting_dir)
# print('results with best model: ', best_model.evaluate(x_test_rgb, y_test_vectorized, batch_size=1, verbose=1))
reconstructed_imgs = new_model.predict(x, batch_size=4, verbose=1)

# x_test_rgb = np.zeros((img_size, img_size, ste, 3))
reconstructed_imgs = reconstructed_imgs.reshape(-1, img_size, img_size)
print(' \n')

# plt.figure(figsize=(n, 3))

display_results_RGB(x, y, reconstructed_imgs, img_size, tests_number, figure_name='best_model_large_dataset')

cv2.namedWindow('test')
cv2.waitKey(1000)
print('finished script!')
cv2.destroyAllWindows()