from __future__ import print_function

import os
import sys

import numpy as np
import tensorflow as tf
from keras import backend as K
from keras.callbacks import TensorBoard
from keras.optimizers import Adam

from data import load_data
from data import oversample
from net import dice_coef
from net import dice_coef_loss
from net import unet

train_images_path = "./data/train/"
valid_images_path = "./data/valid/"
init_weights_path = "./weights_128.h5"
weights_path = "."
log_path = "."

gpu = "0"

epochs = 128
batch_size = 32
base_lr = 1e-5


def train():
    imgs_train, imgs_mask_train, _ = load_data(train_images_path)

    mean = np.mean(imgs_train)
    std = np.std(imgs_train)

    imgs_train -= mean
    imgs_train /= std

    imgs_valid, imgs_mask_valid, _ = load_data(valid_images_path)

    imgs_valid -= mean
    imgs_valid /= std

    imgs_train, imgs_mask_train = oversample(imgs_train, imgs_mask_train)

    model = unet()
    if os.path.exists(init_weights_path):
        model.load_weights(init_weights_path)

    optimizer = Adam(lr=base_lr)
    model.compile(optimizer=optimizer, loss=dice_coef_loss, metrics=[dice_coef])

    if not os.path.exists(log_path):
        os.mkdir(log_path)

    training_log = TensorBoard(log_dir=log_path)

    model.fit(
        imgs_train,
        imgs_mask_train,
        validation_data=(imgs_valid, imgs_mask_valid),
        batch_size=batch_size,
        epochs=epochs,
        shuffle=True,
        callbacks=[training_log],
    )

    if not os.path.exists(weights_path):
        os.mkdir(weights_path)
    model.save_weights(os.path.join(weights_path, "weights_{}.h5".format(epochs)))


if __name__ == "__main__":
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    config.allow_soft_placement = True
    sess = tf.Session(config=config)
    K.set_session(sess)

    if len(sys.argv) > 1:
        gpu = sys.argv[1]
    device = "/gpu:" + gpu

    with tf.device(device):
        train()