from __future__ import print_function

# import packages
from model import unet_model_3d
from keras.utils import plot_model
from keras import callbacks
from keras.callbacks import ModelCheckpoint, CSVLogger, LearningRateScheduler, ReduceLROnPlateau, EarlyStopping

# import load data
from data_handling import load_train_data, load_validatation_data

# import configurations
import configs

# init configs
patch_size = configs.PATCH_SIZE
batch_size = configs.BATCH_SIZE

config = dict()
config["pool_size"] = (2, 2, 2)  # pool size for the max pooling operations
config["image_shape"] = (256, 128, 256)  # This determines what shape the images will be cropped/resampled to.
config["input_shape"] = (patch_size, patch_size, patch_size, 1)  # switch to None to train on the whole image (64, 64, 64) (64, 64, 64)
config["n_labels"] = 4
config["all_modalities"] = ['t1']#]["t1", "t1Gd", "flair", "t2"]
config["training_modalities"] = config["all_modalities"]  # change this if you want to only use some of the modalities
config["nb_channels"] = len(config["training_modalities"])
config["deconvolution"] = False  # if False, will use upsampling instead of deconvolution
config["batch_size"] = batch_size
config["n_epochs"] = 500  # cutoff the training after this many epochs
config["patience"] = 10  # learning rate will be reduced after this many epochs if the validation loss is not improving
config["early_stop"] = 30  # training will be stopped after this many epochs without the validation loss improving
config["initial_learning_rate"] = 0.00005
config["depth"] = configs.DEPTH
config["learning_rate_drop"] = 0.5

image_type = '3d_patches'

# resume training
def resume():
    print('-'*30)
    print('Loading and preprocessing train data...')
    print('-'*30)
    imgs_train, imgs_gtruth_train = load_train_data()
    
    print('-'*30)
    print('Loading and preprocessing validation data...')
    print('-'*30)
    imgs_val, imgs_gtruth_val  = load_validatation_data()
    
    print('-'*30)
    print('Creating and compiling model...')
    print('-'*30)

   # create a model
    model = unet_model_3d(input_shape=config["input_shape"],
                                depth=config["depth"],
                                pool_size=config["pool_size"],
                                n_labels=config["n_labels"],
                                initial_learning_rate=config["initial_learning_rate"],
                                deconvolution=config["deconvolution"])

    model.summary()
    
    checkpoint_filepath_best = 'outputs/' + 'best_weights_125extract_depth5_patch32_88_943_935.h5'
    checkpoint_filepath_best = 'outputs/' + 'best_weights_10extract_depth5_patch32_855_945_935.h5'
    checkpoint_filepath_best = 'outputs/' + 'best_weights_12extract_depth5_patch32_85_946_931_norm_tuned10.h5'
    checkpoint_filepath_best = 'outputs/' + 'best_weights.h5'
    #checkpoint_filepath_best = 'outputs/' + 'best_weights_11extract_depth5_patch32_855_945_935_tunedfrom10extract.h5'
    
    #checkpoint_filepath_best = 'outputs/' + 'best_weights_125extract_depth4_patch32_864_941_932.h5'
    model.load_weights(checkpoint_filepath_best)
    
    print('*'*50)
    print('Load model: ', checkpoint_filepath_best)
    print('*'*50)
    
    #summarize layers
    #print(model.summary())
    # plot graph
    #plot_model(model, to_file='3d_unet.png')
    
    print('-'*30)
    print('Fitting model...')
    print('-'*30)
    
    #============================================================================
    print('training starting..')
    log_filename = 'outputs/' + image_type +'_model_train.csv' 
    
    
    csv_log = callbacks.CSVLogger(log_filename, separator=',', append=True)
    
#    early_stopping = callbacks.EarlyStopping(monitor='val_loss', min_delta=0, patience=5, verbose=0, mode='min')
    
    #checkpoint_filepath = 'outputs/' + image_type +"_best_weight_model_{epoch:03d}_{val_loss:.4f}.hdf5"
    checkpoint_filepath = 'outputs/' + 'weights.h5'
    
    checkpoint = callbacks.ModelCheckpoint(checkpoint_filepath, monitor='val_loss', verbose=1, save_best_only=True, mode='min')
    
    callbacks_list = [csv_log, checkpoint]
    callbacks_list.append(ReduceLROnPlateau(factor=config["learning_rate_drop"], patience=config["patience"],
                                           verbose=True))
    callbacks_list.append(EarlyStopping(verbose=True, patience=config["early_stop"]))

    #============================================================================
    hist = model.fit(imgs_train, imgs_gtruth_train, batch_size=config["batch_size"], nb_epoch=config["n_epochs"], verbose=1, validation_data=(imgs_val,imgs_gtruth_val), shuffle=True, callbacks=callbacks_list) #              validation_split=0.2,
        
     
    model_name = 'outputs/' + image_type + '_model_last'
    model.save(model_name)  # creates a HDF5 file 'my_model.h5'

    
if __name__ == '__main__':
    resume()