import glob

import os

import keras.backend as K

import numpy as np

from keras import Input, Model

from keras.engine.network import Network

from keras.layers import Activation, Concatenate

from keras.optimizers import Adam

from tqdm import tqdm

import fetal_net

import fetal_net.metrics

import fetal_net.preprocess

from fetal.config_utils import get_config

from fetal.utils import get_last_model_path, create_data_file, set_gpu_mem_growth

from fetal_net.data import open_data_file

from fetal_net.generator import get_training_and_validation_generators

from fetal_net.model.fetal_net import fetal_envelope_model

set_gpu_mem_growth()

config = get_config()

if not "dis_model_name" in config:

    config["dis_model_name"] = "discriminator_image"

if not "dis_loss" in config:

    config["dis_loss"] = "binary_crossentropy_loss"

if not "gen_steps" in config:

    config["gen_steps"] = 1

if not "dis_steps" in config:

    config["dis_steps"] = 1

if not "gd_loss_ratio" in config:

    config["gd_loss_ratio"] = 10

class Scheduler:

    def __init__(self, n_itrs_per_epoch_d, n_itrs_per_epoch_g, init_lr, lr_decay, lr_patience):

        self.init_dsteps = n_itrs_per_epoch_d

        self.init_gsteps = n_itrs_per_epoch_g

        self.init_lr = init_lr

        self.lr_decay = lr_decay

        self.lr_patience = lr_patience

        self.dsteps = self.init_dsteps

        self.gsteps = self.init_gsteps

        self.lr = self.init_lr

        self.steps_stuck = 0

        self.best_loss = np.inf

    def get_dsteps(self):

        return self.dsteps

    def get_gsteps(self):

        return self.gsteps

    def get_lr(self):

        return self.lr

    def update_steps(self, n_round, loss):

        if loss < self.best_loss:

            self.steps_stuck = 0

            self.best_loss = loss

        else:

            self.steps_stuck += 1

        if self.steps_stuck > self.lr_patience:

            self.lr *= self.lr_decay

            self.steps_stuck = 0

            print('Reducing LR to {}'.format(self.lr))

        # if key in self.schedules['step_decay']:

        # self.dsteps = max(int(self.init_dsteps * self.schedules['step_decay'][key]), 1)

        # self.gsteps = max(int(self.init_gsteps * self.schedules['step_decay'][key]), 1)

def build_dsc(out_labels, outs):

    s = ''

    for l, o in zip(out_labels, outs):

        s = s + '{}={:.3f}, '.format(l, o)

    return s[:-2] + '|'

def add_noise_to_segs(segs):

    if np.random.choice([True, False]):

        segs = segs.astype(np.float32)

        segs += np.random.normal(0, 0.025, segs.shape)

        segs *= np.random.normal(1, 0.025, segs.shape)

        segs = np.clip(segs, a_min=0, a_max=1)

    return segs

def mul_merge_maps(r, s):

    r_plus = r * s

    r_minus = r * (1 - s)

    return np.concatenate((r_plus, r_minus), axis=1)

def input2discriminator(real_patches, real_segs, semi_patches, semi_segs, d_out_shape, mul_merge=True):

    if mul_merge:

        real = mul_merge_maps(real_patches, add_noise_to_segs(real_segs))

        fake = mul_merge_maps(semi_patches, semi_segs)

    else:

        real = np.concatenate((real_patches, add_noise_to_segs(real_segs)), axis=1)

        fake = np.concatenate((semi_patches, semi_segs), axis=1)

    d_x_batch = np.concatenate((real, fake), axis=0)

    # real : 1, fake : 0

    d_y_batch = np.clip(np.random.uniform(0.9, 1.0, size=[d_x_batch.shape[0]] + list(d_out_shape)[1:]),

                        a_min=0, a_max=1)

    d_y_batch[real.shape[0]:, ...] = 1 - d_y_batch[real.shape[0]:, ...]

    return d_x_batch, d_y_batch

def input2gan(real_patches, real_segs, semi_patches, d_out_shape):

    g_x_batch = [real_patches, semi_patches]

    # set 1 to all labels (real : 1, fake : 0)

    g_y_batch = [

        real_segs,

        np.clip(np.random.uniform(0.9, 1.0, size=[real_patches.shape[0]] + list(d_out_shape)[1:]), a_min=0, a_max=1)

    ]

    return g_x_batch, g_y_batch

def main(overwrite=False):

    # convert input images into an hdf5 file

    if overwrite or not os.path.exists(config["data_file"]):

        create_data_file(config)

    data_file_opened = open_data_file(config["data_file"])

    seg_loss_func = getattr(fetal_net.metrics, config['loss'])

    dis_loss_func = getattr(fetal_net.metrics, config['dis_loss'])

    # instantiate new model

    seg_model_func = getattr(fetal_net.model, config['model_name'])

    gen_model = seg_model_func(input_shape=config["input_shape"],

                               initial_learning_rate=config["initial_learning_rate"],

                               **{'dropout_rate': config['dropout_rate'],

                                  'loss_function': seg_loss_func,

                                  'mask_shape': None if config["weight_mask"] is None else config[

                                      "input_shape"],

                                  'old_model_path': config['old_model']})

    dis_model_func = getattr(fetal_net.model, config['dis_model_name'])

    dis_model = dis_model_func(

        input_shape=[config["input_shape"][0] + config["n_labels"]] + config["input_shape"][1:],

        initial_learning_rate=config["initial_learning_rate"],

        **{'dropout_rate': config['dropout_rate'],

           'loss_function': dis_loss_func})

    if not overwrite \

            and len(glob.glob(config["model_file"] + 'g_*.h5')) > 0:

        # dis_model_path = get_last_model_path(config["model_file"] + 'dis_')

        gen_model_path = get_last_model_path(config["model_file"] + 'g_')

        # print('Loading dis model from: {}'.format(dis_model_path))

        print('Loading gen model from: {}'.format(gen_model_path))

        # dis_model = load_old_model(dis_model_path)

        # gen_model = load_old_model(gen_model_path)

        # dis_model.load_weights(dis_model_path)

        gen_model.load_weights(gen_model_path)

    gen_model.summary()

    dis_model.summary()

    # Build "frozen discriminator"

    frozen_dis_model = Network(

        dis_model.inputs,

        dis_model.outputs,

        name='frozen_discriminator'

    )

    frozen_dis_model.trainable = False

    inputs_real = Input(shape=config["input_shape"])

    inputs_fake = Input(shape=config["input_shape"])

    segs_real = Activation(None, name='seg_real')(gen_model(inputs_real))

    segs_fake = Activation(None, name='seg_fake')(gen_model(inputs_fake))

    valid = Activation(None, name='dis')(frozen_dis_model(Concatenate(axis=1)([segs_fake, inputs_fake])))

    combined_model = Model(inputs=[inputs_real, inputs_fake],

                           outputs=[segs_real, valid])

    combined_model.compile(loss=[seg_loss_func, 'binary_crossentropy'],

                           loss_weights=[1, config["gd_loss_ratio"]],

                           optimizer=Adam(config["initial_learning_rate"]))

    combined_model.summary()

    # get training and testing generators

    train_generator, validation_generator, n_train_steps, n_validation_steps = get_training_and_validation_generators(

        data_file_opened,

        batch_size=config["batch_size"],

        data_split=config["validation_split"],

        overwrite=overwrite,

        validation_keys_file=config["validation_file"],

        training_keys_file=config["training_file"],

        test_keys_file=config["test_file"],

        n_labels=config["n_labels"],

        labels=config["labels"],

        patch_shape=(*config["patch_shape"], config["patch_depth"]),

        validation_batch_size=config["validation_batch_size"],

        augment=config["augment"],

        skip_blank_train=config["skip_blank_train"],

        skip_blank_val=config["skip_blank_val"],

        truth_index=config["truth_index"],

        truth_size=config["truth_size"],

        prev_truth_index=config["prev_truth_index"],

        prev_truth_size=config["prev_truth_size"],

        truth_downsample=config["truth_downsample"],

        truth_crop=config["truth_crop"],

        patches_per_epoch=config["patches_per_epoch"],

        categorical=config["categorical"], is3d=config["3D"],

        drop_easy_patches_train=config["drop_easy_patches_train"],

        drop_easy_patches_val=config["drop_easy_patches_val"])

    # get training and testing generators

    _, semi_generator, _, _ = get_training_and_validation_generators(

        data_file_opened,

        batch_size=config["batch_size"],

        data_split=config["validation_split"],

        overwrite=overwrite,

        validation_keys_file=config["validation_file"],

        training_keys_file=config["training_file"],

        test_keys_file=config["test_file"],

        n_labels=config["n_labels"],

        labels=config["labels"],

        patch_shape=(*config["patch_shape"], config["patch_depth"]),

        validation_batch_size=config["validation_batch_size"],

        val_augment=config["augment"],

        skip_blank_train=config["skip_blank_train"],

        skip_blank_val=config["skip_blank_val"],

        truth_index=config["truth_index"],

        truth_size=config["truth_size"],

        prev_truth_index=config["prev_truth_index"],

        prev_truth_size=config["prev_truth_size"],

        truth_downsample=config["truth_downsample"],

        truth_crop=config["truth_crop"],

        patches_per_epoch=config["patches_per_epoch"],

        categorical=config["categorical"], is3d=config["3D"],

        drop_easy_patches_train=config["drop_easy_patches_train"],

        drop_easy_patches_val=config["drop_easy_patches_val"])

    # start training

    scheduler = Scheduler(config["dis_steps"], config["gen_steps"],

                          init_lr=config["initial_learning_rate"],

                          lr_patience=config["patience"],

                          lr_decay=config["learning_rate_drop"])

    best_loss = np.inf

    for epoch in range(config["n_epochs"]):

        postfix = {'g': None, 'd': None}  # , 'val_g': None, 'val_d': None}

        with tqdm(range(n_train_steps // config["gen_steps"]), dynamic_ncols=True,

                  postfix={'gen': None, 'dis': None, 'val_gen': None, 'val_dis': None, None: None}) as pbar:

            for n_round in pbar:

                # train D

                outputs = np.zeros(dis_model.metrics_names.__len__())

                for i in range(scheduler.get_dsteps()):

                    real_patches, real_segs = next(train_generator)

                    semi_patches, _ = next(semi_generator)

                    d_x_batch, d_y_batch = input2discriminator(real_patches, real_segs,

                                                               semi_patches,

                                                               gen_model.predict(semi_patches,

                                                                                 batch_size=config["batch_size"]),

                                                               dis_model.output_shape)

                    outputs += dis_model.train_on_batch(d_x_batch, d_y_batch)

                if scheduler.get_dsteps():

                    outputs /= scheduler.get_dsteps()

                    postfix['d'] = build_dsc(dis_model.metrics_names, outputs)

                    pbar.set_postfix(**postfix)

                # train G (freeze discriminator)

                outputs = np.zeros(combined_model.metrics_names.__len__())

                for i in range(scheduler.get_gsteps()):

                    real_patches, real_segs = next(train_generator)

                    semi_patches, _ = next(validation_generator)

                    g_x_batch, g_y_batch = input2gan(real_patches, real_segs, semi_patches, dis_model.output_shape)

                    outputs += combined_model.train_on_batch(g_x_batch, g_y_batch)

                outputs /= scheduler.get_gsteps()

                postfix['g'] = build_dsc(combined_model.metrics_names, outputs)

                pbar.set_postfix(**postfix)

            # evaluate on validation set

            dis_metrics = np.zeros(dis_model.metrics_names.__len__(), dtype=float)

            gen_metrics = np.zeros(gen_model.metrics_names.__len__(), dtype=float)

            evaluation_rounds = n_validation_steps

            for n_round in range(evaluation_rounds):  # rounds_for_evaluation:

                val_patches, val_segs = next(validation_generator)

                # D

                if scheduler.get_dsteps() > 0:

                    d_x_test, d_y_test = input2discriminator(val_patches, val_segs,

                                                             val_patches,

                                                             gen_model.predict(val_patches,

                                                                               batch_size=config[

                                                                                   "validation_batch_size"]),

                                                             dis_model.output_shape)

                    dis_metrics += dis_model.evaluate(d_x_test, d_y_test, batch_size=config["validation_batch_size"],

                                                      verbose=0)

                # G

                # gen_x_test, gen_y_test = input2gan(val_patches, val_segs, dis_model.output_shape)

                gen_metrics += gen_model.evaluate(val_patches, val_segs,

                                                  batch_size=config["validation_batch_size"],

                                                  verbose=0)

            dis_metrics /= float(evaluation_rounds)

            gen_metrics /= float(evaluation_rounds)

            # save the model and weights with the best validation loss

            if gen_metrics[0] < best_loss:

                best_loss = gen_metrics[0]

                print('Saving Model...')

                with open(os.path.join(config["base_dir"], "g_{}_{:.3f}.json".format(epoch, gen_metrics[0])),

                          'w') as f:

                    f.write(gen_model.to_json())

                gen_model.save_weights(

                    os.path.join(config["base_dir"], "g_{}_{:.3f}.h5".format(epoch, gen_metrics[0])))

            postfix['val_d'] = build_dsc(dis_model.metrics_names, dis_metrics)

            postfix['val_g'] = build_dsc(gen_model.metrics_names, gen_metrics)

            # pbar.set_postfix(**postfix)

            print('val_d: ' + postfix['val_d'], end=' | ')

            print('val_g: ' + postfix['val_g'])

            # pbar.refresh()

            # update step sizes, learning rates

            scheduler.update_steps(epoch, gen_metrics[0])

            K.set_value(dis_model.optimizer.lr, scheduler.get_lr())

            K.set_value(combined_model.optimizer.lr, scheduler.get_lr())

    data_file_opened.close()

if __name__ == "__main__":

    main(overwrite=config["overwrite"])