"""Contains custom loss, dice coefficient, and optimizer classes."""

import tensorflow as tf

class DiceVAELoss(object):

    """Implements custom dice-VAE loss."""

    def __init__(self,

                 name='custom_loss',

                 data_format='channels_last',

                 **kwargs):

        self.axis = (0, 1, 2, 3) if data_format == 'channels_last' else (0, 2, 3, 4)

    def __call__(self, x, y, y_pred, y_vae, z_mean, z_logvar, sample_weight=None):

        l2_loss = tf.reduce_mean((x - y_vae) ** 2)

        kld_loss = tf.reduce_mean(z_mean ** 2 + tf.math.exp(z_logvar) - z_logvar - 1.0)

        # Calculate dice loss.

        intersection = tf.reduce_sum(y_pred * y, axis=self.axis)

        pred = tf.reduce_sum(y_pred ** 2, axis=self.axis)

        true = tf.reduce_sum(y ** 2, axis=self.axis)

        dice_loss = tf.reduce_mean(1.0 - (2.0 * intersection + 1.0) / (pred + true + 1.0))

        return dice_loss + 0.1*l2_loss + 0.1*kld_loss

class DiceCoefficient(object):

    """Implements dice coefficient for binary classification."""

    def __init__(self,

                 name='dice_coefficient',

                 data_format='channels_last'):

        self.name = name

        self.data_format = data_format

    def __call__(self, y_true, y_pred):

        dice_axes = (0, 1, 2) if self.data_format == 'channels_last' else (0, 2, 3, 4)

        onehot_axis = -1 if self.data_format == 'channels_last' else 1

        # Mask out values that correspond to values < 0.5.

        mask = tf.reduce_max(y_pred, axis=onehot_axis, keepdims=True)

        mask = tf.cast(mask > 0.5, tf.float32)

        # Create one-hot encoding of predictions.

        out_ch = y_pred.shape[onehot_axis]

        y_pred = tf.argmax(y_pred, axis=onehot_axis, output_type=tf.int32)

        y_pred = tf.one_hot(y_pred, out_ch, axis=onehot_axis, dtype=tf.float32)

        y_pred *= mask

        # Compute dice score.

        intersection = tf.reduce_sum(y_pred * y_true, axis=dice_axes)

        pred = tf.reduce_sum(y_pred, axis=dice_axes)

        true = tf.reduce_sum(y_true, axis=dice_axes)

        macroavg = tf.reduce_mean((2.0 * intersection + 1.0) / (pred + true + 1.0))

        microavg = tf.reduce_sum(y_pred * y_true) / (tf.reduce_sum(y_pred) + tf.reduce_sum(y_true))

        return macroavg, microavg

class ScheduledOptim(tf.keras.optimizers.Adam):

    """Adam optimizer that allows for scheduling every epoch."""

    def __init__(self,

                 learning_rate=1e-4,

                 beta_1=0.9,

                 beta_2=0.999,

                 epsilon=1e-7,

                 amsgrad=False,

                 name='Adam',

                 n_epochs=300,

                 **kwargs):

        super(ScheduledOptim, self).__init__(

                                        learning_rate=learning_rate,

                                        beta_1=beta_1,

                                        beta_2=beta_2,

                                        epsilon=epsilon,

                                        amsgrad=amsgrad,

                                        name=name,

                                        **kwargs)

        self.init_lr = tf.constant(learning_rate, dtype=tf.float32)

        self.n_epochs = float(n_epochs)

    def __call__(self, epoch):

        new_lr = self.init_lr * ((1.0 - epoch / self.n_epochs) ** 0.9)

        self._set_hyper('learning_rate', new_lr)