import numpy as np

import random

import json

from glob import glob

from keras.models import model_from_json,load_model

from keras.preprocessing.image import ImageDataGenerator

from keras.callbacks import  ModelCheckpoint,Callback,LearningRateScheduler

import keras.backend as K

from model import Unet_model

from losses import *

#from keras.utils.visualize_util import plot

class SGDLearningRateTracker(Callback):

    def on_epoch_begin(self, epoch, logs={}):

        optimizer = self.model.optimizer

        lr = K.get_value(optimizer.lr)

        decay = K.get_value(optimizer.decay)

        lr=lr/10

        decay=decay*10

        K.set_value(optimizer.lr, lr)

        K.set_value(optimizer.decay, decay)

        print('LR changed to:',lr)

        print('Decay changed to:',decay)

class Training(object):

    def __init__(self, batch_size,nb_epoch,load_model_resume_training=None):

        self.batch_size = batch_size

        self.nb_epoch = nb_epoch

        #loading model from path to resume previous training without recompiling the whole model

        if load_model_resume_training is not None:

            self.model =load_model(load_model_resume_training,custom_objects={'gen_dice_loss': gen_dice_loss,'dice_whole_metric':dice_whole_metric,'dice_core_metric':dice_core_metric,'dice_en_metric':dice_en_metric})

            print("pre-trained model loaded!")

        else:

            unet =Unet_model(img_shape=(128,128,4))

            self.model=unet.model

            print("U-net CNN compiled!")

    def fit_unet(self,X33_train,Y_train,X_patches_valid=None,Y_labels_valid=None):

        train_generator=self.img_msk_gen(X33_train,Y_train,9999)

        checkpointer = ModelCheckpoint(filepath='brain_segmentation/ResUnet.{epoch:02d}_{val_loss:.3f}.hdf5', verbose=1)

        self.model.fit_generator(train_generator,steps_per_epoch=len(X33_train)//self.batch_size,epochs=self.nb_epoch, validation_data=(X_patches_valid,Y_labels_valid),verbose=1, callbacks = [checkpointer,SGDLearningRateTracker()])

        #self.model.fit(X33_train,Y_train, epochs=self.nb_epoch,batch_size=self.batch_size,validation_data=(X_patches_valid,Y_labels_valid),verbose=1, callbacks = [checkpointer,SGDLearningRateTracker()])

    def img_msk_gen(self,X33_train,Y_train,seed):

        '''

        a custom generator that performs data augmentation on both patches and their corresponding targets (masks)

        '''

        datagen = ImageDataGenerator(horizontal_flip=True,data_format="channels_last")

        datagen_msk = ImageDataGenerator(horizontal_flip=True,data_format="channels_last")

        image_generator = datagen.flow(X33_train,batch_size=4,seed=seed)

        y_generator = datagen_msk.flow(Y_train,batch_size=4,seed=seed)

        while True:

            yield(image_generator.next(), y_generator.next())

    def save_model(self, model_name):

        '''

        INPUT string 'model_name': path where to save model and weights, without extension

        Saves current model as json and weights as h5df file

        '''

        model_tosave = '{}.json'.format(model_name)

        weights = '{}.hdf5'.format(model_name)

        json_string = self.model.to_json()

        self.model.save_weights(weights)

        with open(model_tosave, 'w') as f:

            json.dump(json_string, f)

        print ('Model saved.')

    def load_model(self, model_name):

        '''

        Load a model

        INPUT  (1) string 'model_name': filepath to model and weights, not including extension

        OUTPUT: Model with loaded weights. can fit on model using loaded_model=True in fit_model method

        '''

        print ('Loading model {}'.format(model_name))

        model_toload = '{}.json'.format(model_name)

        weights = '{}.hdf5'.format(model_name)

        with open(model_toload) as f:

            m = next(f)

        model_comp = model_from_json(json.loads(m))

        model_comp.load_weights(weights)

        print ('Model loaded.')

        self.model = model_comp

        return model_comp

if __name__ == "__main__":

    #set arguments

    #reload already trained model to resume training

    model_to_load="Models/ResUnet.04_0.646.hdf5" 

    #save=None

    #compile the model

    brain_seg = Training(batch_size=4,nb_epoch=3,load_model_resume_training=model_to_load)

    print("number of trainabale parameters:",brain_seg.model.count_params())

    #print(brain_seg.model.summary())

    #plot(brain_seg.model, to_file='model_architecture.png', show_shapes=True)

    #load data from disk

    Y_labels=np.load("y_training.npy").astype(np.uint8)

    X_patches=np.load("x_training.npy").astype(np.float32)

    Y_labels_valid=np.load("y_valid.npy").astype(np.uint8)

    X_patches_valid=np.load("x_valid.npy").astype(np.float32)

    print("loading patches done\n")

    # fit model

    brain_seg.fit_unet(X_patches,Y_labels,X_patches_valid,Y_labels_valid)#*

    #if save is not None:

    #    brain_seg.save_model('models/' + save)