import argparse
from keras import optimizers
import matplotlib.pyplot as plt
import numpy as np
import datetime
from keras.callbacks import TensorBoard
import glob
import os
import tensorflow as tf
from models import *
from utils.lr_controller import ReduceLROnPlateau
from utils.data_loader import data_loader, data_loader_multi_channel
from utils.utils import img_comp
from utils.loss import loss_mse_ssim

parser = argparse.ArgumentParser()
parser.add_argument("--gpu_id", type=int, default=1)
parser.add_argument("--gpu_memory_fraction", type=float, default=0.3)
parser.add_argument("--mixed_precision_training", type=int, default=1)
parser.add_argument("--data_dir", type=str, default="../dataset/train/F-actin")
parser.add_argument("--save_weights_dir", type=str, default="../trained_models")
parser.add_argument("--model_name", type=str, default="DFCAN")
parser.add_argument("--patch_height", type=int, default=128)
parser.add_argument("--patch_width", type=int, default=128)
parser.add_argument("--input_channels", type=int, default=9)
parser.add_argument("--scale_factor", type=int, default=2)
parser.add_argument("--norm_flag", type=int, default=1)
parser.add_argument("--iterations", type=int, default=1000000)
parser.add_argument("--sample_interval", type=int, default=1000)
parser.add_argument("--validate_interval", type=int, default=2000)
parser.add_argument("--validate_num", type=int, default=500)
parser.add_argument("--batch_size", type=int, default=4)
parser.add_argument("--start_lr", type=float, default=1e-4)
parser.add_argument("--lr_decay_factor", type=float, default=0.5)
parser.add_argument("--load_weights", type=int, default=0)
parser.add_argument("--optimizer_name", type=str, default="adam")

args = parser.parse_args()
gpu_id = str(args.gpu_id)
gpu_memory_fraction = args.gpu_memory_fraction
mixed_precision_training = str(args.mixed_precision_training)
data_dir = args.data_dir
save_weights_dir = args.save_weights_dir
validate_interval = args.validate_interval
batch_size = args.batch_size
start_lr = args.start_lr
lr_decay_factor = args.lr_decay_factor
patch_height = args.patch_height
patch_width = args.patch_width
input_channels = args.input_channels
scale_factor = args.scale_factor
norm_flag = args.norm_flag
validate_num = args.validate_num
iterations = args.iterations
load_weights = args.load_weights
optimizer_name = args.optimizer_name
model_name = args.model_name
sample_interval = args.sample_interval

os.environ["TF_ENABLE_AUTO_MIXED_PRECISION"] = mixed_precision_training
os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=gpu_memory_fraction)
tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

data_name = data_dir.split('/')[-1]
if input_channels == 1:
    save_weights_name = model_name + '-SISR_' + data_name
    cur_data_loader = data_loader
    train_images_path = data_dir + '/training_wf/'
    validate_images_path = data_dir + '/validate_wf/'
else:
    save_weights_name = model_name + '-SIM_' + data_name
    cur_data_loader = data_loader_multi_channel
    train_images_path = data_dir + '/training/'
    validate_images_path = data_dir + '/validate/'
save_weights_path = save_weights_dir + '/' + save_weights_name + '/'
train_gt_path = data_dir + '/training_gt/'
validate_gt_path = data_dir + '/validate_gt/'
sample_path = save_weights_path + 'sampled_img/'

if not os.path.exists(save_weights_path):
    os.mkdir(save_weights_path)
if not os.path.exists(sample_path):
    os.mkdir(sample_path)

# --------------------------------------------------------------------------------
#                           select models and optimizer
# --------------------------------------------------------------------------------
modelFns = {'DFCAN': DFCAN16.DFCAN}
modelFN = modelFns[model_name]
optimizer_g = optimizers.adam(lr=start_lr, beta_1=0.9, beta_2=0.999)

# --------------------------------------------------------------------------------
#                              define combined model
# --------------------------------------------------------------------------------
g = modelFN((patch_height, patch_width, input_channels))
g.compile(loss=loss_mse_ssim, optimizer=optimizer_g)
lr_controller = ReduceLROnPlateau(model=g, factor=lr_decay_factor, patience=10, mode='min', min_delta=1e-4,
                                  cooldown=0, min_lr=start_lr * 0.1, verbose=1)

# --------------------------------------------------------------------------------
#                                 about Tensorboard
# --------------------------------------------------------------------------------
log_path = save_weights_path + 'graph'
if not os.path.exists(log_path):
    os.mkdir(log_path)
callback = TensorBoard(log_path)
callback.set_model(g)
train_names = 'training_loss'
val_names = ['val_MSE', 'val_SSIM', 'val_PSNR', 'val_NRMSE']


def write_log(callback, names, logs, batch_no):
    summary = tf.Summary()
    summary_value = summary.value.add()
    summary_value.simple_value = logs
    summary_value.tag = names
    callback.writer.add_summary(summary, batch_no)
    callback.writer.flush()


# --------------------------------------------------------------------------------
#                             Sample and validate
# --------------------------------------------------------------------------------
def Validate(iter, sample=0):
    validate_path = glob.glob(validate_images_path + '*')
    validate_path.sort()
    if sample == 1:
        r, c = 3, 3
        mses, nrmses, psnrs, ssims = [], [], [], []
        img_show, gt_show, output_show = [], [], []
        validate_path = np.random.choice(validate_path, size=r)
        for path in validate_path:
            [img, gt] = cur_data_loader([path], validate_images_path, validate_gt_path, patch_height,
                                        patch_width, 1, norm_flag=norm_flag, scale=scale_factor)
            output = np.squeeze(g.predict(img))
            mses, nrmses, psnrs, ssims = img_comp(gt, output, mses, nrmses, psnrs, ssims)
            img_show.append(np.squeeze(np.mean(img, 3)))
            gt_show.append(np.squeeze(gt))
            output_show.append(output)
            # show some examples
        fig, axs = plt.subplots(r, c)
        cnt = 0
        for row in range(r):
            axs[row, 1].set_title('MSE=%.4f, SSIM=%.4f, PSNR=%.4f' % (mses[row], ssims[row], psnrs[row]))
            for col, image in enumerate([img_show, output_show, gt_show]):
                axs[row, col].imshow(np.squeeze(image[row]))
                axs[row, col].axis('off')
            cnt += 1
        fig.savefig(sample_path + '%d.png' % iter)
        plt.close()
    else:
        if validate_num < validate_path.__len__():
            validate_path = validate_path[0:validate_num]
        mses, nrmses, psnrs, ssims = [], [], [], []
        for path in validate_path:
            [img, gt] = cur_data_loader([path], validate_images_path, validate_gt_path, patch_height,
                                        patch_width, 1, norm_flag=norm_flag, scale=scale_factor)
            output = np.squeeze(g.predict(img))
            mses, nrmses, psnrs, ssims = img_comp(gt, output, mses, nrmses, psnrs, ssims)

        # if best, save weights.best
        g.save_weights(save_weights_path + 'weights.latest')
        if min(validate_nrmse) > np.mean(nrmses):
            g.save_weights(save_weights_path + 'weights.best')

        validate_nrmse.append(np.mean(nrmses))
        curlr = lr_controller.on_epoch_end(iter, np.mean(nrmses))
        write_log(callback, val_names[0], np.mean(mses), iter)
        write_log(callback, val_names[1], np.mean(ssims), iter)
        write_log(callback, val_names[2], np.mean(psnrs), iter)
        write_log(callback, val_names[3], np.mean(nrmses), iter)
        write_log(callback, 'lr', curlr, iter)


# --------------------------------------------------------------------------------
#                             if exist, load weights
# --------------------------------------------------------------------------------
if load_weights:
    if os.path.exists(save_weights_path + 'weights.best'):
        g.save_weights(save_weights_path + 'weights.best')
        print('Loading weights successfully: ' + save_weights_path + 'weights.best')
    elif os.path.exists(save_weights_path + 'weights.latest'):
        g.save_weights(save_weights_path + 'weights.latest')
        print('Loading weights successfully: ' + save_weights_path + 'weights.latest')


# --------------------------------------------------------------------------------
#                                    training
# --------------------------------------------------------------------------------
start_time = datetime.datetime.now()
loss_record = []
validate_nrmse = [np.Inf]
lr_controller.on_train_begin()
images_path = glob.glob(train_images_path + '/*')
for it in range(iterations):
    # ------------------------------------
    #         train generator
    # ------------------------------------
    input_g, gt_g = cur_data_loader(images_path, train_images_path, train_gt_path, patch_height, patch_width,
                                    batch_size, norm_flag=norm_flag, scale=scale_factor)
    loss_generator = g.train_on_batch(input_g, gt_g)
    loss_record.append(loss_generator)

    elapsed_time = datetime.datetime.now() - start_time
    print("%d epoch: time: %s, g_loss = %s" % (it + 1, elapsed_time, loss_generator))

    if (it + 1) % sample_interval == 0:
        images_path = glob.glob(train_images_path + '/*')
        Validate(it + 1, sample=1)

    if (it + 1) % validate_interval == 0:
        Validate(it + 1, sample=0)
        write_log(callback, train_names, np.mean(loss_record), it + 1)
        loss_record = []

[((16, 9, 16, 34), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((62, 14, 62, 80), 'tensorflow.GPUOptions', 'tf.GPUOptions', (), '', True, 'import tensorflow as tf\n'), ((91, 14, 91, 68), 'keras.optimizers.adam', 'optimizers.adam', (), '', False, 'from keras import optimizers\n'), ((98, 16, 99, 79), 'utils.lr_controller.ReduceLROnPlateau', 'ReduceLROnPlateau', (), '', False, 'from utils.lr_controller import ReduceLROnPlateau\n'), ((107, 11, 107, 32), 'keras.callbacks.TensorBoard', 'TensorBoard', ({(107, 23, 107, 31): 'log_path'}, {}), '(log_path)', False, 'from keras.callbacks import TensorBoard\n'), ((191, 13, 191, 36), 'datetime.datetime.now', 'datetime.datetime.now', ({}, {}), '()', False, 'import datetime\n'), ((195, 14, 195, 49), 'glob.glob', 'glob.glob', ({(195, 24, 195, 48): "train_images_path + '/*'"}, {}), "(train_images_path + '/*')", False, 'import glob\n'), ((81, 7, 81, 40), 'os.path.exists', 'os.path.exists', ({(81, 22, 81, 39): 'save_weights_path'}, {}), '(save_weights_path)', False, 'import os\n'), ((82, 4, 82, 31), 'os.mkdir', 'os.mkdir', ({(82, 13, 82, 30): 'save_weights_path'}, {}), '(save_weights_path)', False, 'import os\n'), ((83, 7, 83, 34), 'os.path.exists', 'os.path.exists', ({(83, 22, 83, 33): 'sample_path'}, {}), '(sample_path)', False, 'import os\n'), ((84, 4, 84, 25), 'os.mkdir', 'os.mkdir', ({(84, 13, 84, 24): 'sample_path'}, {}), '(sample_path)', False, 'import os\n'), ((105, 7, 105, 31), 'os.path.exists', 'os.path.exists', ({(105, 22, 105, 30): 'log_path'}, {}), '(log_path)', False, 'import os\n'), ((106, 4, 106, 22), 'os.mkdir', 'os.mkdir', ({(106, 13, 106, 21): 'log_path'}, {}), '(log_path)', False, 'import os\n'), ((114, 14, 114, 26), 'tensorflow.Summary', 'tf.Summary', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((126, 20, 126, 57), 'glob.glob', 'glob.glob', ({(126, 30, 126, 56): "validate_images_path + '*'"}, {}), "(validate_images_path + '*')", False, 'import glob\n'), ((180, 7, 180, 57), 'os.path.exists', 'os.path.exists', ({(180, 22, 180, 56): "(save_weights_path + 'weights.best')"}, {}), "(save_weights_path + 'weights.best')", False, 'import os\n'), ((63, 18, 63, 57), 'tensorflow.ConfigProto', 'tf.ConfigProto', (), '', True, 'import tensorflow as tf\n'), ((132, 24, 132, 63), 'numpy.random.choice', 'np.random.choice', (), '', True, 'import numpy as np\n'), ((142, 19, 142, 37), 'matplotlib.pyplot.subplots', 'plt.subplots', ({(142, 32, 142, 33): 'r', (142, 35, 142, 36): 'c'}, {}), '(r, c)', True, 'import matplotlib.pyplot as plt\n'), ((151, 8, 151, 19), 'matplotlib.pyplot.close', 'plt.close', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((183, 9, 183, 61), 'os.path.exists', 'os.path.exists', ({(183, 24, 183, 60): "(save_weights_path + 'weights.latest')"}, {}), "(save_weights_path + 'weights.latest')", False, 'import os\n'), ((205, 19, 205, 42), 'datetime.datetime.now', 'datetime.datetime.now', ({}, {}), '()', False, 'import datetime\n'), ((209, 22, 209, 57), 'glob.glob', 'glob.glob', ({(209, 32, 209, 56): "train_images_path + '/*'"}, {}), "(train_images_path + '/*')", False, 'import glob\n'), ((137, 41, 137, 89), 'utils.utils.img_comp', 'img_comp', ({(137, 50, 137, 52): 'gt', (137, 54, 137, 60): 'output', (137, 62, 137, 66): 'mses', (137, 68, 137, 74): 'nrmses', (137, 76, 137, 81): 'psnrs', (137, 83, 137, 88): 'ssims'}, {}), '(gt, output, mses, nrmses, psnrs, ssims)', False, 'from utils.utils import img_comp\n'), ((160, 41, 160, 89), 'utils.utils.img_comp', 'img_comp', ({(160, 50, 160, 52): 'gt', (160, 54, 160, 60): 'output', (160, 62, 160, 66): 'mses', (160, 68, 160, 74): 'nrmses', (160, 76, 160, 81): 'psnrs', (160, 83, 160, 88): 'ssims'}, {}), '(gt, output, mses, nrmses, psnrs, ssims)', False, 'from utils.utils import img_comp\n'), ((164, 33, 164, 48), 'numpy.mean', 'np.mean', ({(164, 41, 164, 47): 'nrmses'}, {}), '(nrmses)', True, 'import numpy as np\n'), ((167, 30, 167, 45), 'numpy.mean', 'np.mean', ({(167, 38, 167, 44): 'nrmses'}, {}), '(nrmses)', True, 'import numpy as np\n'), ((168, 49, 168, 64), 'numpy.mean', 'np.mean', ({(168, 57, 168, 63): 'nrmses'}, {}), '(nrmses)', True, 'import numpy as np\n'), ((169, 42, 169, 55), 'numpy.mean', 'np.mean', ({(169, 50, 169, 54): 'mses'}, {}), '(mses)', True, 'import numpy as np\n'), ((170, 42, 170, 56), 'numpy.mean', 'np.mean', ({(170, 50, 170, 55): 'ssims'}, {}), '(ssims)', True, 'import numpy as np\n'), ((171, 42, 171, 56), 'numpy.mean', 'np.mean', ({(171, 50, 171, 55): 'psnrs'}, {}), '(psnrs)', True, 'import numpy as np\n'), ((172, 42, 172, 57), 'numpy.mean', 'np.mean', ({(172, 50, 172, 56): 'nrmses'}, {}), '(nrmses)', True, 'import numpy as np\n'), ((214, 41, 214, 61), 'numpy.mean', 'np.mean', ({(214, 49, 214, 60): 'loss_record'}, {}), '(loss_record)', True, 'import numpy as np\n'), ((139, 27, 139, 41), 'numpy.squeeze', 'np.squeeze', ({(139, 38, 139, 40): 'gt'}, {}), '(gt)', True, 'import numpy as np\n'), ((138, 39, 138, 54), 'numpy.mean', 'np.mean', ({(138, 47, 138, 50): 'img', (138, 52, 138, 53): '(3)'}, {}), '(img, 3)', True, 'import numpy as np\n'), ((147, 37, 147, 59), 'numpy.squeeze', 'np.squeeze', ({(147, 48, 147, 58): 'image[row]'}, {}), '(image[row])', True, 'import numpy as np\n')]