def aspcap_mask(elem, dr=None):
    """
    | To load ASPCAP elements window masks
    | DR14 Elements: ``'C', 'CI', 'N', 'O', 'Na', 'Mg', 'Al', 'Si', 'P', 'S', 'K', 'Ca', 'TI', 'TiII', 'V', 'Cr', 'Mn',
                     'Fe', 'Co', 'Ni', 'Cu', 'Ge', 'Ce', 'Rb', 'Y', 'Nd'``

    :param elem: element name
    :type elem: str
    :param dr: apogee dr
    :type dr: int
    :return: mask
    :rtype: ndarray[bool]
    :History: 2018-Mar-24 - Written - Henry Leung (University of Toronto)
    """
    if elem.lower() == 'c1':
        elem = 'CI'

    elif elem.lower() == 'ti2':
        elem = 'TiII'

    dr = apogee_default_dr(dr=dr)

    if 14 <= dr <= 16:
        aspcap_code = 'l31c'
        elem_list = ['C', 'CI', 'N', 'O', 'Na', 'Mg', 'Al', 'Si', 'P', 'S', 'K', 'Ca', 'TI', 'TiII', 'V', 'Cr', 'Mn',
                     'Fe', 'Co', 'Ni', 'Cu', 'Ge', 'Ce', 'Rb', 'Y', 'Nd']
    else:
        raise ValueError('Only DR14-DR16 is supported currently')

    masks = np.load(
        os.path.join(astroNN.data.datapath(), f'aspcap_{aspcap_code}_masks.npy'))

    try:
        # turn everything to lowercase to avoid case-related issue
        index = [x.lower() for x in elem_list].index(elem.lower())
    except ValueError:
        # nicely handle if element not found
        print(f'Element not found, the only elements for dr{dr} supported are {elem_list}')
        return None

    return [(masks & 2 ** index) != 0][0] class BayesianCNNBase(NeuralNetMaster, ABC):
    """
    Top-level class for a Bayesian convolutional neural network

    :History: 2018-Jan-06 - Written - Henry Leung (University of Toronto)
    """

    def __init__(self):
        super().__init__()
        self.name = 'Bayesian Convolutional Neural Network'
        self._model_type = 'BCNN'
        self.initializer = None
        self.activation = None
        self._last_layer_activation = None
        self.num_filters = None
        self.filter_len = None
        self.pool_length = None
        self.num_hidden = None
        self.reduce_lr_epsilon = None
        self.reduce_lr_min = None
        self.reduce_lr_patience = None
        self.l1 = None
        self.l2 = None
        self.maxnorm = None
        self.inv_model_precision = None  # inverse model precision
        self.dropout_rate = 0.2
        self.length_scale = 3  # prior length scale
        self.mc_num = 100  # increased to 100 due to high performance VI on GPU implemented on 14 April 2018 (Henry)
        self.val_size = 0.1
        self.disable_dropout = False

        self.input_norm_mode = 1
        self.labels_norm_mode = 2

        self.keras_model_predict = None

    def pre_training_checklist_child(self, input_data, labels, sample_weights):
        input_data, labels = self.pre_training_checklist_master(input_data, labels)

        # check if exists (existing means the model has already been trained (e.g. fine-tuning), so we do not need calculate mean/std again)
        if self.input_normalizer is None:
            self.input_normalizer = Normalizer(mode=self.input_norm_mode)
            self.labels_normalizer = Normalizer(mode=self.labels_norm_mode)

            norm_data = self.input_normalizer.normalize(input_data)
            self.input_mean, self.input_std = self.input_normalizer.mean_labels, self.input_normalizer.std_labels
            norm_labels = self.labels_normalizer.normalize(labels)
            self.labels_mean, self.labels_std = self.labels_normalizer.mean_labels, self.labels_normalizer.std_labels
        else:
            norm_data = self.input_normalizer.normalize(input_data, calc=False)
            norm_labels = self.labels_normalizer.normalize(labels, calc=False)

        # No need to care about Magic number as loss function looks for magic num in y_true only
        norm_data.update({"input_err": (input_data['input_err'] / self.input_std['input']),
                          "labels_err": input_data['labels_err'] / self.labels_std['output']})
        norm_labels.update({"variance_output": norm_labels['output']})

        if self.keras_model is None:  # only compile if there is no keras_model, e.g. fine-tuning does not required
            self.compile()

        self.train_idx, self.val_idx = train_test_split(np.arange(self.num_train + self.val_num),
                                                        test_size=self.val_size)

        norm_data_training = {}
        norm_data_val = {}
        norm_labels_training = {}
        norm_labels_val = {}
        for name in norm_data.keys():
            norm_data_training.update({name: norm_data[name][self.train_idx]})
            norm_data_val.update({name: norm_data[name][self.val_idx]})
        for name in norm_labels.keys():
            norm_labels_training.update({name: norm_labels[name][self.train_idx]})
            norm_labels_val.update({name: norm_labels[name][self.val_idx]})
        
        if sample_weights is not None:        
            sample_weights_training = sample_weights[self.train_idx]
            sample_weights_val = sample_weights[self.val_idx]
        else:
            sample_weights_training = None
            sample_weights_val = None

        self.inv_model_precision = (2 * self.num_train * self.l2) / (self.length_scale ** 2 * (1 - self.dropout_rate))

        self.training_generator = BayesianCNNDataGenerator(batch_size=self.batch_size,
                                                           shuffle=True,
                                                           steps_per_epoch=self.num_train // self.batch_size,
                                                           data=[norm_data_training,
                                                                 norm_labels_training],
                                                           manual_reset=False, 
                                                           sample_weights=sample_weights_training)

        val_batchsize = self.batch_size if len(self.val_idx) > self.batch_size else len(self.val_idx)
        self.validation_generator = BayesianCNNDataGenerator(batch_size=val_batchsize,
                                                             shuffle=False,
                                                             steps_per_epoch=max(self.val_num // self.batch_size, 1),
                                                             data=[norm_data_val,
                                                                   norm_labels_val],
                                                             manual_reset=True,
                                                             sample_weights=sample_weights_val)

        return norm_data_training, norm_data_val, norm_labels_training, norm_labels_val, sample_weights_training, sample_weights_val

    def compile(self, optimizer=None,
                loss=None,
                metrics=None,
                weighted_metrics=None,
                loss_weights=None,
                sample_weight_mode=None):
        if optimizer is not None:
            self.optimizer = optimizer
        elif self.optimizer is None or self.optimizer == 'adam':
            self.optimizer = Adam(learning_rate=self.lr, beta_1=self.beta_1, beta_2=self.beta_2, epsilon=self.optimizer_epsilon,
                                  decay=0.0)
        if metrics is not None:
            self.metrics = metrics
        if self.task == 'regression':
            if self._last_layer_activation is None:
                self._last_layer_activation = 'linear'
        elif self.task == 'classification':
            if self._last_layer_activation is None:
                self._last_layer_activation = 'softmax'
        elif self.task == 'binary_classification':
            if self._last_layer_activation is None:
                self._last_layer_activation = 'sigmoid'
        else:
            raise RuntimeError('Only "regression", "classification" and "binary_classification" are supported')

        self.keras_model, self.keras_model_predict, self.output_loss, self.variance_loss = self.model()
        
        if self.task == 'regression':
            self._output_loss = lambda predictive, labelerr: mse_lin_wrapper(predictive, labelerr)
        elif self.task == 'classification':
            self._output_loss = lambda predictive, labelerr: bayesian_categorical_crossentropy_wrapper(predictive)
        elif self.task == 'binary_classification':
            self._output_loss = lambda predictive, labelerr: bayesian_binary_crossentropy_wrapper(predictive)
        else:
            raise RuntimeError('Only "regression", "classification" and "binary_classification" are supported')
        
        # all zero losss as dummy lose
        if self.task == 'regression':
            self.metrics = [mean_absolute_error, mean_error] if not self.metrics else self.metrics
            self.keras_model.compile(optimizer=self.optimizer,
                                     loss=zeros_loss,
                                     metrics=self.metrics,
                                     weighted_metrics=weighted_metrics,
                                     sample_weight_mode=sample_weight_mode)
        elif self.task == 'classification':
            self.metrics = [categorical_accuracy] if not self.metrics else self.metrics
            self.keras_model.compile(optimizer=self.optimizer,
                                     loss=zeros_loss,
                                     metrics={'output': self.metrics},
                                     weighted_metrics=weighted_metrics,
                                     sample_weight_mode=sample_weight_mode)
        elif self.task == 'binary_classification':
            self.metrics = [binary_accuracy] if not self.metrics else self.metrics
            self.keras_model.compile(optimizer=self.optimizer,
                                     loss=zeros_loss,
                                     metrics={'output': self.metrics},
                                     weighted_metrics=weighted_metrics,
                                     sample_weight_mode=sample_weight_mode)
        
        # inject custom training step if needed
        try:
            self.custom_train_step()
        except NotImplementedError:
            pass
        except TypeError:
            self.keras_model.train_step = self.custom_train_step
       # inject custom testing  step if needed
        try:
            self.custom_test_step()
        except NotImplementedError:
            pass
        except TypeError:
            self.keras_model.test_step = self.custom_test_step

        return None

    def custom_train_step(self, data):
        """
        Custom training logic

        :param data:
        :return:
        """
        data = data_adapter.expand_1d(data)
        x, y, sample_weight = data_adapter.unpack_x_y_sample_weight(data)

        # Run forward pass.
        with tf.GradientTape() as tape:
            y_pred = self.keras_model(x, training=True)
            self.keras_model.compiled_loss._losses = self._output_loss(y_pred[1], x['labels_err'])
            self.keras_model.compiled_loss._losses = nest.map_structure(self.keras_model.compiled_loss._get_loss_object, self.keras_model.compiled_loss._losses)
            self.keras_model.compiled_loss._losses = nest.flatten(self.keras_model.compiled_loss._losses)
            loss = self.keras_model.compiled_loss(y, y_pred, sample_weight, regularization_losses=self.keras_model.losses)
            
        # Run backwards pass.
        self.keras_model.optimizer.minimize(loss, self.keras_model.trainable_variables, tape=tape)
        self.keras_model.compiled_metrics.update_state(y, y_pred, sample_weight)
        # Collect metrics to return
        return_metrics = {}
        for metric in self.keras_model.metrics:
            result = metric.result()
            if isinstance(result, dict):
                return_metrics.update(result)
            else:
                return_metrics[metric.name] = result
        return return_metrics

    
    def custom_test_step(self, data):
        data = data_adapter.expand_1d(data)
        x, y, sample_weight = data_adapter.unpack_x_y_sample_weight(data)
                
        y_pred = self.keras_model(x, training=False)
        # Updates stateful loss metrics.
        temploss = self._output_loss(y_pred[1], x['labels_err'])
        self.keras_model.compiled_loss._losses = temploss
        self.keras_model.compiled_loss._losses = nest.map_structure(self.keras_model.compiled_loss._get_loss_object, self.keras_model.compiled_loss._losses)
        self.keras_model.compiled_loss._losses = nest.flatten(self.keras_model.compiled_loss._losses)

        self.keras_model.compiled_loss(
            y, y_pred, sample_weight, regularization_losses=self.keras_model.losses)
        
        self.keras_model.compiled_metrics.update_state(y, y_pred, sample_weight)
        # Collect metrics to return
        return_metrics = {}
            
        for metric in self.keras_model.metrics:
            result = metric.result()
            if isinstance(result, dict):
                return_metrics.update(result)
            else:
                return_metrics[metric.name] = result
        return return_metrics


    def fit(self, input_data, labels, inputs_err=None, labels_err=None, sample_weights=None, experimental=False):
        """
        Train a Bayesian neural network

        :param input_data: Data to be trained with neural network
        :type input_data: ndarray
        :param labels: Labels to be trained with neural network
        :type labels: ndarray
        :param inputs_err: Error for input_data (if any), same shape with input_data.
        :type inputs_err: Union([NoneType, ndarray])
        :param labels_err: Labels error (if any)
        :type labels_err: Union([NoneType, ndarray])
        :param sample_weights: Sample weights (if any)
        :type sample_weights: Union([NoneType, ndarray])
        :return: None
        :rtype: NoneType
        :History:
            | 2018-Jan-06 - Written - Henry Leung (University of Toronto)
            | 2018-Apr-12 - Updated - Henry Leung (University of Toronto)
        """
        if inputs_err is None:
            inputs_err = np.zeros_like(input_data)

        if labels_err is None:
            labels_err = np.zeros_like(labels)

        # TODO: allow named inputs too??
        input_data = {"input": input_data, "input_err": inputs_err, "labels_err": labels_err}
        labels = {"output": labels, "variance_output": labels}

        # Call the checklist to create astroNN folder and save parameters
        norm_data_training, norm_data_val, norm_labels_training, norm_labels_val, sample_weights_training, sample_weights_val = self.pre_training_checklist_child(input_data, labels, sample_weights)
        
        # norm_data_training['labels_err'] = norm_data_training['labels_err'].filled(MAGIC_NUMBER).astype(np.float32)
        
        # TODO: fix the monitor name
        reduce_lr = ReduceLROnPlateau(monitor='val_output_mean_absolute_error', factor=0.5, min_delta=self.reduce_lr_epsilon,
                                      patience=self.reduce_lr_patience, min_lr=self.reduce_lr_min, mode='min',
                                      verbose=2)

        self.virtual_cvslogger = VirutalCSVLogger()

        self.__callbacks = [reduce_lr, self.virtual_cvslogger]  # default must have unchangeable callbacks

        if self.callbacks is not None:
            if isinstance(self.callbacks, list):
                self.__callbacks.extend(self.callbacks)
            else:
                self.__callbacks.append(self.callbacks)

        start_time = time.time()

        if experimental:
            dataset = tf.data.Dataset.from_tensor_slices((norm_data_training, norm_labels_training, sample_weights_training)).batch(self.batch_size).shuffle(5000, reshuffle_each_iteration=True).prefetch(tf.data.AUTOTUNE)
            val_dataset = tf.data.Dataset.from_tensor_slices((norm_data_val, norm_labels_val, sample_weights_val)).batch(self.batch_size).prefetch(tf.data.AUTOTUNE)

            self.history = self.keras_model.fit(dataset,
                                                validation_data=val_dataset,
                                                epochs=self.max_epochs, verbose=self.verbose,
                                                workers=os.cpu_count() // 2,
                                                callbacks=self.__callbacks,
                                                use_multiprocessing=MULTIPROCESS_FLAG)
        else:
            self.history = self.keras_model.fit(self.training_generator,
                                                validation_data=self.validation_generator,
                                                epochs=self.max_epochs, verbose=self.verbose,
                                                workers=os.cpu_count() // 2,
                                                callbacks=self.__callbacks,
                                                use_multiprocessing=MULTIPROCESS_FLAG)

        print(f'Completed Training, {(time.time() - start_time):.{2}f}s in total')
        if self.autosave is True:
            # Call the post training checklist to save parameters
            self.save()

        return None

    def fit_on_batch(self, input_data, labels, inputs_err=None, labels_err=None, sample_weights=None):
        """
        Train a Bayesian neural network by running a single gradient update on all of your data, suitable for fine-tuning

        :param input_data: Data to be trained with neural network
        :type input_data: ndarray
        :param labels: Labels to be trained with neural network
        :type labels: ndarray
        :param inputs_err: Error for input_data (if any), same shape with input_data.
        :type inputs_err: Union([NoneType, ndarray])
        :param labels_err: Labels error (if any)
        :type labels_err: Union([NoneType, ndarray])
        :param sample_weights: Sample weights (if any)
        :type sample_weights: Union([NoneType, ndarray])
        :return: None
        :rtype: NoneType
        :History:
            | 2018-Aug-25 - Written - Henry Leung (University of Toronto)
        """
        self.has_model_check()

        if inputs_err is None:
            inputs_err = np.zeros_like(input_data)

        if labels_err is None:
            labels_err = np.zeros_like(labels)

        input_data = {"input": input_data, "input_err": inputs_err, "labels_err": labels_err}
        labels = {"output": labels, "variance_output": labels}

        # check if exists (existing means the model has already been trained (e.g. fine-tuning), so we do not need calculate mean/std again)
        if self.input_normalizer is None:
            self.input_normalizer = Normalizer(mode=self.input_norm_mode)
            self.labels_normalizer = Normalizer(mode=self.labels_norm_mode)

            norm_data = self.input_normalizer.normalize(input_data)
            self.input_mean, self.input_std = self.input_normalizer.mean_labels, self.input_normalizer.std_labels
            norm_labels = self.labels_normalizer.normalize(labels)
            self.labels_mean, self.labels_std = self.labels_normalizer.mean_labels, self.labels_normalizer.std_labels
        else:
            norm_data = self.input_normalizer.normalize(input_data, calc=False)
            norm_labels = self.labels_normalizer.normalize(labels, calc=False)

        # No need to care about Magic number as loss function looks for magic num in y_true only
        norm_data.update({"input_err": (input_data['input_err'] / self.input_std['input']),
                          "labels_err": input_data['labels_err'] / self.labels_std['output']})
        norm_labels.update({"variance_output": norm_labels['output']})

        start_time = time.time()

        fit_generator = BayesianCNNDataGenerator(batch_size=input_data['input'].shape[0],
                                                 shuffle=False,
                                                 steps_per_epoch=1,
                                                 data=[norm_data,
                                                       norm_labels], 
                                                 sample_weights=sample_weights)

        score = self.keras_model.fit(fit_generator,
                                     epochs=1,
                                     verbose=self.verbose,
                                     workers=os.cpu_count(),
                                     use_multiprocessing=MULTIPROCESS_FLAG)

        print(f'Completed Training on Batch, {(time.time() - start_time):.{2}f}s in total')

        return None

    def post_training_checklist_child(self):
        self.keras_model.save(self.fullfilepath + _astroNN_MODEL_NAME)
        print(_astroNN_MODEL_NAME + f' saved to {(self.fullfilepath + _astroNN_MODEL_NAME)}')

        self.hyper_txt.write(f"Dropout Rate: {self.dropout_rate} \n")
        self.hyper_txt.flush()
        self.hyper_txt.close()

        data = {'id': self.__class__.__name__ if self._model_identifier is None else self._model_identifier,
                'pool_length': self.pool_length,
                'filterlen': self.filter_len,
                'filternum': self.num_filters,
                'hidden': self.num_hidden,
                'input': self._input_shape,
                'labels': self._labels_shape,
                'task': self.task,
                'last_layer_activation': self._last_layer_activation,
                'activation': self.activation,
                'input_mean': dict_np_to_dict_list(self.input_mean),
                'inv_tau': self.inv_model_precision,
                'length_scale': self.length_scale,
                'labels_mean': dict_np_to_dict_list(self.labels_mean),
                'input_std': dict_np_to_dict_list(self.input_std),
                'labels_std': dict_np_to_dict_list(self.labels_std),
                'valsize': self.val_size,
                'targetname': self.targetname,
                'dropout_rate': self.dropout_rate,
                'l1': self.l1,
                'l2': self.l2,
                'maxnorm': self.maxnorm,
                'input_norm_mode': self.input_normalizer.normalization_mode,
                'labels_norm_mode': self.labels_normalizer.normalization_mode,
                'input_names': self.input_names,
                'output_names': self.output_names,
                'batch_size': self.batch_size}

        with open(self.fullfilepath + '/astroNN_model_parameter.json', 'w') as f:
            json.dump(data, f, indent=4, sort_keys=True)

    def predict(self, input_data, inputs_err=None):
        """
        Test model, High performance version designed for fast variational inference on GPU

        :param input_data: Data to be inferred with neural network
        :type input_data: ndarray
        :param inputs_err: Error for input_data, same shape with input_data.
        :type inputs_err: Union([NoneType, ndarray])
        :return: prediction and prediction uncertainty
        :History:
            | 2018-Jan-06 - Written - Henry Leung (University of Toronto)
            | 2018-Apr-12 - Updated - Henry Leung (University of Toronto)
        """
        self.has_model_check()

        if gpu_availability() is False and self.mc_num > 25:
            warnings.warn(f'You are using CPU version Tensorflow, doing {self.mc_num} times Monte Carlo Inference can '
                          f'potentially be very slow! \n '
                          f'A possible fix is to decrease the mc_num parameter of the model to do less MC Inference \n'
                          f'This is just a warning, and will not shown if mc_num < 25 on CPU')
            if self.mc_num < 2:
                raise AttributeError("mc_num cannot be smaller than 2")

        # if no error array then just zeros
        if inputs_err is None:
            inputs_err = np.zeros_like(input_data)
        else:
            inputs_err = np.atleast_2d(inputs_err)
            inputs_err /= self.input_std['input']

        input_data = {"input": input_data, "input_err": inputs_err}
        input_data = self.pre_testing_checklist_master(input_data)

        if self.input_normalizer is not None:
            input_array = self.input_normalizer.normalize(input_data, calc=False)
        else:
            # Prevent shallow copy issue
            input_array = np.array(input_data)
            input_array -= self.input_mean['input']
            input_array /= self.input_std['input']

        total_test_num = input_data['input'].shape[0]  # Number of testing data

        # for number of training data smaller than batch_size
        if total_test_num < self.batch_size:
            batch_size = total_test_num
        else:
            batch_size = self.batch_size

        # Due to the nature of how generator works, no overlapped prediction
        data_gen_shape = (total_test_num // batch_size) * batch_size
        remainder_shape = total_test_num - data_gen_shape  # Remainder from generator

        norm_data_main = {}
        norm_data_remainder = {}
        for name in input_array.keys():
            norm_data_main.update({name: input_array[name][:data_gen_shape]})
            norm_data_remainder.update({name: input_array[name][data_gen_shape:]})

        # Data Generator for prediction
        with tqdm(total=total_test_num, unit="sample") as pbar:
            pbar.set_postfix({'Monte-Carlo': self.mc_num})
            # suppress pfor warning from TF
            old_level = tf.get_logger().level
            tf.get_logger().setLevel('ERROR')

            prediction_generator = BayesianCNNPredDataGenerator(batch_size=batch_size,
                                                                shuffle=False,
                                                                steps_per_epoch=data_gen_shape // batch_size,
                                                                data=[norm_data_main], 
                                                                pbar=pbar)

            new = FastMCInference(self.mc_num)(self.keras_model_predict)
            
            result = np.asarray(new.predict(prediction_generator))

            if remainder_shape != 0:  # deal with remainder
                remainder_generator = BayesianCNNPredDataGenerator(batch_size=remainder_shape,
                                                                    shuffle=False,
                                                                    steps_per_epoch=1,
                                                                    data=[norm_data_remainder], 
                                                                    pbar=pbar)
                remainder_result = np.asarray(new.predict(remainder_generator))
                if remainder_shape == 1:
                    remainder_result = np.expand_dims(remainder_result, axis=0)
                result = np.concatenate((result, remainder_result))
                
            tf.get_logger().setLevel(old_level)

        # in case only 1 test data point, in such case we need to add a dimension
        if result.ndim < 3 and batch_size == 1:
            result = np.expand_dims(result, axis=0)

        half_first_dim = result.shape[1] // 2  # result.shape[1] is guarantee an even number, otherwise sth is wrong

        predictions = result[:, :half_first_dim, 0]  # mean prediction
        mc_dropout_uncertainty = result[:, :half_first_dim, 1] * (self.labels_std['output'] ** 2)  # model uncertainty
        predictions_var = np.exp(result[:, half_first_dim:, 0]) * (
                self.labels_std['output'] ** 2)  # predictive uncertainty

        if self.labels_normalizer is not None:
            predictions = self.labels_normalizer.denormalize(
                list_to_dict([self.keras_model.output_names[0]], predictions))
            predictions = predictions['output']
        else:
            predictions *= self.labels_std['output']
            predictions += self.labels_mean['output']

        if self.task == 'regression':
            # Predictive variance
            pred_var = predictions_var + mc_dropout_uncertainty  # epistemic plus aleatoric uncertainty
            pred_uncertainty = np.sqrt(pred_var)  # Convert back to std error

            # final correction from variance to standard derivation
            mc_dropout_uncertainty = np.sqrt(mc_dropout_uncertainty)
            predictive_uncertainty = np.sqrt(predictions_var)

        elif self.task == 'classification':
            # we want entropy for classification uncertainty
            predicted_class = np.argmax(predictions, axis=1)
            mc_dropout_uncertainty = np.ones_like(predicted_class, dtype=float)
            predictive_uncertainty = np.ones_like(predicted_class, dtype=float)

            # center variance
            predictions_var -= 1.
            for i in range(predicted_class.shape[0]):
                all_prediction = np.array(predictions[i, :])
                mc_dropout_uncertainty[i] = - np.sum(all_prediction * np.log(all_prediction))
                predictive_uncertainty[i] = predictions_var[i, predicted_class[i]]

            pred_uncertainty = mc_dropout_uncertainty + predictive_uncertainty
            # We only want the predicted class back
            predictions = predicted_class

        elif self.task == 'binary_classification':
            # we want entropy for classification uncertainty, so need prediction in logits space
            mc_dropout_uncertainty = - np.sum(predictions * np.log(predictions), axis=0)
            # need to activate before round to int so that the prediction is always 0 or 1
            predictions = np.rint(sigmoid(predictions))
            predictive_uncertainty = predictions_var
            pred_uncertainty = mc_dropout_uncertainty + predictions_var

        else:
            raise AttributeError('Unknown Task')

        return predictions, {'total': pred_uncertainty, 'model': mc_dropout_uncertainty,
                             'predictive': predictive_uncertainty}


    def predict_dataset(self, file):
        class BayesianCNNPredDataGeneratorV2(GeneratorMaster):
            def __init__(self, batch_size, shuffle, steps_per_epoch, manual_reset=False, pbar=None, nn_model=None):
                super().__init__(batch_size=batch_size, shuffle=shuffle, steps_per_epoch=steps_per_epoch, data=None,
                                manual_reset=manual_reset)
                self.pbar = pbar

                # initial idx
                self.idx_list = self._get_exploration_order(range(len(file)))
                self.current_idx = 0
                self.nn_model = nn_model

            def _data_generation(self, idx_list_temp):
                # Generate data
                inputs = self.nn_model.input_normalizer.normalize({"input": file[idx_list_temp], "input_err": np.zeros_like(file[idx_list_temp])}, calc=False)
                x = self.input_d_checking(inputs, np.arange(len(idx_list_temp)))
                return x

            def __getitem__(self, index):
                x = self._data_generation(self.idx_list[index * self.batch_size: (index + 1) * self.batch_size])
                if self.pbar: self.pbar.update(self.batch_size)
                return x

            def on_epoch_end(self):
                # shuffle the list when epoch ends for the next epoch
                self.idx_list = self._get_exploration_order(range(len(file)))
            
        self.has_model_check()

        if gpu_availability() is False and self.mc_num > 25:
            warnings.warn(f'You are using CPU version Tensorflow, doing {self.mc_num} times Monte Carlo Inference can '
                          f'potentially be very slow! \n '
                          f'A possible fix is to decrease the mc_num parameter of the model to do less MC Inference \n'
                          f'This is just a warning, and will not shown if mc_num < 25 on CPU')
            if self.mc_num < 2:
                raise AttributeError("mc_num cannot be smaller than 2")
        
        total_test_num = len(file)  # Number of testing data

        # for number of training data smaller than batch_size
        if total_test_num < self.batch_size:
            batch_size = total_test_num
        else:
            batch_size = self.batch_size

        # Due to the nature of how generator works, no overlapped prediction
        data_gen_shape = (total_test_num // batch_size) * batch_size
        remainder_shape = total_test_num - data_gen_shape  # Remainder from generator

        # Data Generator for prediction
        with tqdm(total=total_test_num, unit="sample") as pbar:
            pbar.set_postfix({'Monte-Carlo': self.mc_num})
            # suppress pfor warning from TF
            old_level = tf.get_logger().level
            tf.get_logger().setLevel('ERROR')
            prediction_generator = BayesianCNNPredDataGeneratorV2(batch_size=batch_size,
                                                                shuffle=False,
                                                                steps_per_epoch=data_gen_shape // batch_size,
                                                                pbar=pbar, 
                                                                nn_model=self)

            new = FastMCInference(self.mc_num)(self.keras_model_predict)
            
            result = np.asarray(new.predict(prediction_generator))

            if remainder_shape != 0:  # deal with remainder
                remainder_generator = BayesianCNNPredDataGeneratorV2(batch_size=remainder_shape,
                                                                    shuffle=False,
                                                                    steps_per_epoch=1,
                                                                    pbar=pbar, 
                                                                    nn_model=self)
                remainder_result = np.asarray(new.predict(remainder_generator))
                if remainder_shape == 1:
                    remainder_result = np.expand_dims(remainder_result, axis=0)
                result = np.concatenate((result, remainder_result))
                
            tf.get_logger().setLevel(old_level)

        # in case only 1 test data point, in such case we need to add a dimension
        if result.ndim < 3 and batch_size == 1:
            result = np.expand_dims(result, axis=0)

        half_first_dim = result.shape[1] // 2  # result.shape[1] is guarantee an even number, otherwise sth is wrong

        predictions = result[:, :half_first_dim, 0]  # mean prediction
        mc_dropout_uncertainty = result[:, :half_first_dim, 1] * (self.labels_std['output'] ** 2)  # model uncertainty
        predictions_var = np.exp(result[:, half_first_dim:, 0]) * (
                self.labels_std['output'] ** 2)  # predictive uncertainty

        if self.labels_normalizer is not None:
            predictions = self.labels_normalizer.denormalize(
                list_to_dict([self.keras_model.output_names[0]], predictions))
            predictions = predictions['output']
        else:
            predictions *= self.labels_std['output']
            predictions += self.labels_mean['output']

        if self.task == 'regression':
            # Predictive variance
            pred_var = predictions_var + mc_dropout_uncertainty  # epistemic plus aleatoric uncertainty
            pred_uncertainty = np.sqrt(pred_var)  # Convert back to std error

            # final correction from variance to standard derivation
            mc_dropout_uncertainty = np.sqrt(mc_dropout_uncertainty)
            predictive_uncertainty = np.sqrt(predictions_var)

        elif self.task == 'classification':
            # we want entropy for classification uncertainty
            predicted_class = np.argmax(predictions, axis=1)
            mc_dropout_uncertainty = np.ones_like(predicted_class, dtype=float)
            predictive_uncertainty = np.ones_like(predicted_class, dtype=float)

            # center variance
            predictions_var -= 1.
            for i in range(predicted_class.shape[0]):
                all_prediction = np.array(predictions[i, :])
                mc_dropout_uncertainty[i] = - np.sum(all_prediction * np.log(all_prediction))
                predictive_uncertainty[i] = predictions_var[i, predicted_class[i]]

            pred_uncertainty = mc_dropout_uncertainty + predictive_uncertainty
            # We only want the predicted class back
            predictions = predicted_class

        elif self.task == 'binary_classification':
            # we want entropy for classification uncertainty, so need prediction in logits space
            mc_dropout_uncertainty = - np.sum(predictions * np.log(predictions), axis=0)
            # need to activate before round to int so that the prediction is always 0 or 1
            predictions = np.rint(sigmoid(predictions))
            predictive_uncertainty = predictions_var
            pred_uncertainty = mc_dropout_uncertainty + predictions_var

        else:
            raise AttributeError('Unknown Task')

        return predictions, {'total': pred_uncertainty, 'model': mc_dropout_uncertainty,
                             'predictive': predictive_uncertainty}
        
    def evaluate(self, input_data, labels, inputs_err=None, labels_err=None):
        """
        Evaluate neural network by provided input data and labels and get back a metrics score

        :param input_data: Data to be trained with neural network
        :type input_data: ndarray
        :param labels: Labels to be trained with neural network
        :type labels: ndarray
        :param inputs_err: Error for input_data (if any), same shape with input_data.
        :type inputs_err: Union([NoneType, ndarray])
        :param labels_err: Labels error (if any)
        :type labels_err: Union([NoneType, ndarray])
        :return: metrics score dictionary
        :rtype: dict
        :History: 2018-May-20 - Written - Henry Leung (University of Toronto)
        """
        self.has_model_check()

        if inputs_err is None:
            inputs_err = np.zeros_like(input_data)

        if labels_err is None:
            labels_err = np.zeros_like(labels)

        input_data = {"input": input_data}
        labels = {"output": labels}

        # check if exists (existing means the model has already been trained (e.g. fine-tuning), so we do not need calculate mean/std again)
        if self.input_normalizer is None:
            self.input_normalizer = Normalizer(mode=self.input_norm_mode)
            self.labels_normalizer = Normalizer(mode=self.labels_norm_mode)

            norm_data = self.input_normalizer.normalize(input_data)
            self.input_mean, self.input_std = self.input_normalizer.mean_labels, self.input_normalizer.std_labels
            norm_labels = self.labels_normalizer.normalize(labels)
            self.labels_mean, self.labels_std = self.labels_normalizer.mean_labels, self.labels_normalizer.std_labels
        else:
            norm_data = self.input_normalizer.normalize(input_data, calc=False)
            norm_labels = self.labels_normalizer.normalize(labels, calc=False)

        # No need to care about Magic number as loss function looks for magic num in y_true only
        norm_input_err = inputs_err / self.input_std['input']
        norm_labels_err = labels_err / self.labels_std['output']

        norm_data.update({"input_err": norm_input_err, "labels_err": norm_labels_err})
        norm_labels.update({"variance_output": norm_labels["output"]})

        total_num = input_data['input'].shape[0]
        eval_batchsize = self.batch_size if total_num > self.batch_size else total_num
        steps = total_num // self.batch_size if total_num > self.batch_size else 1

        start_time = time.time()
        print("Starting Evaluation")
        
        # suppress pfor warning from TF
        old_level = tf.get_logger().level
        tf.get_logger().setLevel('ERROR')

        evaluate_generator = BayesianCNNDataGenerator(batch_size=eval_batchsize,
                                                      shuffle=False,
                                                      steps_per_epoch=steps,
                                                      data=[norm_data,
                                                            norm_labels])

        scores = self.keras_model.evaluate(evaluate_generator)
        
        tf.get_logger().setLevel(old_level)

        if isinstance(scores, float):  # make sure scores is iterable
            scores = list(str(scores))
        outputname = self.keras_model.output_names
        funcname = self.keras_model.metrics_names

        print(f'Completed Evaluation, {(time.time() - start_time):.{2}f}s elapsed')

        return list_to_dict(funcname, scores)
    
    @deprecated_copy_signature(fit)
    def train(self, *args, **kwargs):
        return self.fit(*args, **kwargs)

    @deprecated_copy_signature(fit_on_batch)
    def train_on_batch(self, *args, **kwargs):
        return self.fit_on_batch(*args, **kwargs)
    
    @deprecated_copy_signature(predict)
    def test(self, *args, **kwargs):
        return self.predict(*args, **kwargs) class CNNBase(NeuralNetMaster, ABC):
    """Top-level class for a convolutional neural network"""

    def __init__(self):
        """
        NAME:
            __init__
        PURPOSE:
            To define astroNN convolutional neural network
        HISTORY:
            2018-Jan-06 - Written - Henry Leung (University of Toronto)
        """
        super().__init__()
        self.name = 'Convolutional Neural Network'
        self._model_type = 'CNN'
        self._model_identifier = None
        self.initializer = None
        self.activation = None
        self._last_layer_activation = None
        self.num_filters = None
        self.filter_len = None
        self.pool_length = None
        self.num_hidden = None
        self.reduce_lr_epsilon = None
        self.reduce_lr_min = None
        self.reduce_lr_patience = None
        self.l1 = None
        self.l2 = None
        self.maxnorm = None
        self.dropout_rate = 0.0
        self.val_size = 0.1
        self.early_stopping_min_delta = 0.0001
        self.early_stopping_patience = 4

        self.input_norm_mode = 1
        self.labels_norm_mode = 2

    def compile(self, optimizer=None,
                loss=None,
                metrics=None,
                weighted_metrics=None,
                loss_weights=None,
                sample_weight_mode=None):

        if optimizer is not None:
            self.optimizer = optimizer
        elif self.optimizer is None or self.optimizer == 'adam':
            self.optimizer = Adam(learning_rate=self.lr, beta_1=self.beta_1, beta_2=self.beta_2, epsilon=self.optimizer_epsilon,
                                  decay=0.0)
        if metrics is not None:
            self.metrics = metrics

        if self.task == 'regression':
            self._last_layer_activation = 'linear'
            loss_func = mean_squared_error if not loss else loss
            self.metrics = [mean_absolute_error, mean_error] if not self.metrics else self.metrics
        elif self.task == 'classification':
            self._last_layer_activation = 'softmax'
            loss_func = categorical_crossentropy if not loss else loss
            self.metrics = [categorical_accuracy] if not self.metrics else self.metrics
        elif self.task == 'binary_classification':
            self._last_layer_activation = 'sigmoid'
            loss_func = binary_crossentropy if not loss else loss
            self.metrics = [binary_accuracy] if not self.metrics else self.metrics
        else:
            raise RuntimeError('Only "regression", "classification" and "binary_classification" are supported')

        self.keras_model = self.model()

        self.keras_model.compile(loss=loss_func,
                                 optimizer=self.optimizer,
                                 metrics=self.metrics,
                                 weighted_metrics=weighted_metrics,
                                 loss_weights=loss_weights,
                                 sample_weight_mode=sample_weight_mode)

        # inject custom training step if needed
        try:
            self.custom_train_step()
        except NotImplementedError:
            pass
        except TypeError:
            self.keras_model.train_step = self.custom_train_step
       # inject custom testing  step if needed
        try:
            self.custom_test_step()
        except NotImplementedError:
            pass
        except TypeError:
            self.keras_model.test_step = self.custom_test_step

        return None

    def pre_training_checklist_child(self, input_data, labels, sample_weights):
        # on top of checklist, convert input_data/labels to dict
        input_data, labels = self.pre_training_checklist_master(input_data, labels)

        # check if exists (existing means the model has already been trained (e.g. fine-tuning)
        # so we do not need calculate mean/std again)
        if self.input_normalizer is None:
            self.input_normalizer = Normalizer(mode=self.input_norm_mode)
            self.labels_normalizer = Normalizer(mode=self.labels_norm_mode)
            norm_data = self.input_normalizer.normalize(input_data)
            self.input_mean, self.input_std = self.input_normalizer.mean_labels, self.input_normalizer.std_labels
            norm_labels = self.labels_normalizer.normalize(labels)
            self.labels_mean, self.labels_std = self.labels_normalizer.mean_labels, self.labels_normalizer.std_labels
        else:
            norm_data = self.input_normalizer.normalize(input_data, calc=False)
            norm_labels = self.labels_normalizer.normalize(labels, calc=False)
        if self.keras_model is None:  # only compile if there is no keras_model, e.g. fine-tuning does not required
            self.compile()
        
        norm_data = self._tensor_dict_sanitize(norm_data, self.keras_model.input_names)
        norm_labels = self._tensor_dict_sanitize(norm_labels, self.keras_model.output_names)

        self.train_idx, self.val_idx = train_test_split(np.arange(self.num_train + self.val_num),
                                                        test_size=self.val_size)

        norm_data_training = {}
        norm_data_val = {}
        norm_labels_training = {}
        norm_labels_val = {}
        for name in norm_data.keys():
            norm_data_training.update({name: norm_data[name][self.train_idx]})
            norm_data_val.update({name: norm_data[name][self.val_idx]})
        for name in norm_labels.keys():
            norm_labels_training.update({name: norm_labels[name][self.train_idx]})
            norm_labels_val.update({name: norm_labels[name][self.val_idx]})

        if sample_weights is not None:        
            sample_weights_training = sample_weights[self.train_idx]
            sample_weights_val = sample_weights[self.val_idx]
        else:
            sample_weights_training = None
            sample_weights_val = None

        self.training_generator = CNNDataGenerator(
            batch_size=self.batch_size,
            shuffle=True,
            steps_per_epoch=self.num_train // self.batch_size,
            data=[norm_data_training, norm_labels_training],
            manual_reset=False, 
            sample_weights=sample_weights_training)

        val_batchsize = self.batch_size if len(self.val_idx) > self.batch_size else len(self.val_idx)
        self.validation_generator = CNNDataGenerator(
            batch_size=val_batchsize,
            shuffle=False,
            steps_per_epoch=max(self.val_num // self.batch_size, 1),
            data=[norm_data_val, norm_labels_val],
            manual_reset=True, 
            sample_weights=sample_weights_val)

        return input_data, labels

    def fit(self, input_data, labels, sample_weights=None):
        """
        Train a Convolutional neural network

        :param input_data: Data to be trained with neural network
        :type input_data: ndarray
        :param labels: Labels to be trained with neural network
        :type labels: ndarray
        :param sample_weights: Sample weights (if any)
        :type sample_weights: Union([NoneType, ndarray])
        :return: None
        :rtype: NoneType
        :History: 2017-Dec-06 - Written - Henry Leung (University of Toronto)
        """
        # Call the checklist to create astroNN folder and save parameters
        self.pre_training_checklist_child(input_data, labels, sample_weights)

        reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.5,
                                      min_delta=self.reduce_lr_epsilon,
                                      patience=self.reduce_lr_patience, min_lr=self.reduce_lr_min, mode='min',
                                      verbose=2)

        early_stopping = EarlyStopping(monitor='val_loss', min_delta=self.early_stopping_min_delta,
                                       patience=self.early_stopping_patience, verbose=2, mode='min')

        self.virtual_cvslogger = VirutalCSVLogger()

        self.__callbacks = [reduce_lr, self.virtual_cvslogger]  # default must have unchangeable callbacks

        if self.callbacks is not None:
            if isinstance(self.callbacks, list):
                self.__callbacks.extend(self.callbacks)
            else:
                self.__callbacks.append(self.callbacks)

        start_time = time.time()

        self.history = self.keras_model.fit(x=self.training_generator,
                                            validation_data=self.validation_generator,
                                            epochs=self.max_epochs, verbose=self.verbose,
                                            workers=os.cpu_count(),
                                            callbacks=self.__callbacks,
                                            use_multiprocessing=MULTIPROCESS_FLAG)

        print(f'Completed Training, {(time.time() - start_time):.{2}f}s in total')

        if self.autosave is True:
            # Call the post training checklist to save parameters
            self.save()

        return None

    def fit_on_batch(self, input_data, labels, sample_weights=None):
        """
        Train a neural network by running a single gradient update on all of your data, suitable for fine-tuning

        :param input_data: Data to be trained with neural network
        :type input_data: ndarray
        :param labels: Labels to be trained with neural network
        :type labels: ndarray
        :param sample_weights: Sample weights (if any)
        :type sample_weights: Union([NoneType, ndarray])
        :return: None
        :rtype: NoneType
        :History: 2018-Aug-22 - Written - Henry Leung (University of Toronto)
        """

        input_data, labels = self.pre_training_checklist_master(input_data, labels)

        # check if exists (existing means the model has already been trained (e.g. fine-tuning),
        # so we do not need calculate mean/std again)
        if self.input_normalizer is None:
            self.input_normalizer = Normalizer(mode=self.input_norm_mode)
            self.labels_normalizer = Normalizer(mode=self.labels_norm_mode)

            norm_data = self.input_normalizer.normalize(input_data)
            self.input_mean, self.input_std = self.input_normalizer.mean_labels, self.input_normalizer.std_labels
            norm_labels = self.labels_normalizer.normalize(labels)
            self.labels_mean, self.labels_std = self.labels_normalizer.mean_labels, self.labels_normalizer.std_labels
        else:
            norm_data = self.input_normalizer.normalize(input_data, calc=False)
            norm_labels = self.labels_normalizer.normalize(labels, calc=False)

        start_time = time.time()

        fit_generator = CNNDataGenerator(batch_size=input_data['input'].shape[0],
                                         shuffle=False,
                                         steps_per_epoch=1,
                                         data=[norm_data, norm_labels], 
                                         sample_weights=sample_weights)

        scores = self.keras_model.fit(x=fit_generator,
                                      epochs=1,
                                      verbose=self.verbose,
                                      workers=os.cpu_count(),
                                      use_multiprocessing=MULTIPROCESS_FLAG)

        print(f'Completed Training on Batch, {(time.time() - start_time):.{2}f}s in total')

        return None

    def post_training_checklist_child(self):
        self.keras_model.save(self.fullfilepath + _astroNN_MODEL_NAME)
        print(_astroNN_MODEL_NAME + f' saved to {(self.fullfilepath + _astroNN_MODEL_NAME)}')

        self.hyper_txt.write(f"Dropout Rate: {self.dropout_rate} \n")
        self.hyper_txt.flush()
        self.hyper_txt.close()

        data = {'id': self.__class__.__name__ if self._model_identifier is None else self._model_identifier,
                'pool_length': self.pool_length,
                'filterlen': self.filter_len,
                'filternum': self.num_filters,
                'hidden': self.num_hidden,
                'input': self._input_shape,
                'labels': self._labels_shape,
                'task': self.task,
                'last_layer_activation': self._last_layer_activation,
                'activation': self.activation,
                'input_mean': dict_np_to_dict_list(self.input_mean),
                'labels_mean': dict_np_to_dict_list(self.labels_mean),
                'input_std': dict_np_to_dict_list(self.input_std),
                'labels_std': dict_np_to_dict_list(self.labels_std),
                'valsize': self.val_size,
                'targetname': self.targetname,
                'dropout_rate': self.dropout_rate,
                'l1': self.l1,
                'l2': self.l2,
                'maxnorm': self.maxnorm,
                'input_norm_mode': self.input_normalizer.normalization_mode,
                'labels_norm_mode': self.labels_normalizer.normalization_mode,
                'input_names': self.input_names,
                'output_names': self.output_names,
                'batch_size': self.batch_size}

        with open(self.fullfilepath + '/astroNN_model_parameter.json', 'w') as f:
            json.dump(data, f, indent=4, sort_keys=True)

    def predict(self, input_data):
        """
        Use the neural network to do inference

        :param input_data: Data to be inferred with neural network
        :type input_data: ndarray
        :return: prediction and prediction uncertainty
        :rtype: ndarry
        :History: 2017-Dec-06 - Written - Henry Leung (University of Toronto)
        """
        self.has_model_check()
        input_data = self.pre_testing_checklist_master(input_data)

        input_array = self.input_normalizer.normalize(input_data, calc=False)
        total_test_num = input_data['input'].shape[0]  # Number of testing data

        # for number of training data smaller than batch_size
        if total_test_num < self.batch_size:
            self.batch_size = total_test_num

        # Due to the nature of how generator works, no overlapped prediction
        data_gen_shape = (total_test_num // self.batch_size) * self.batch_size
        remainder_shape = total_test_num - data_gen_shape  # Remainder from generator

        # TODO: named output????
        predictions = np.zeros((total_test_num, self._labels_shape['output']))

        norm_data_main = {}
        norm_data_remainder = {}
        for name in input_array.keys():
            norm_data_main.update({name: input_array[name][:data_gen_shape]})
            norm_data_remainder.update({name: input_array[name][data_gen_shape:]})
            
        norm_data_main = self._tensor_dict_sanitize(norm_data_main, self.keras_model.input_names)
        norm_data_remainder = self._tensor_dict_sanitize(norm_data_remainder, self.keras_model.input_names)

        # Data Generator for prediction
        with tqdm(total=total_test_num, unit="sample") as pbar:
            prediction_generator = CNNPredDataGenerator(batch_size=self.batch_size,
                                                        shuffle=False,
                                                        steps_per_epoch=total_test_num // self.batch_size,
                                                        data=[norm_data_main],
                                                        pbar=pbar)
            predictions[:data_gen_shape] = np.asarray(self.keras_model.predict(prediction_generator))

            if remainder_shape != 0:
                remainder_generator = CNNPredDataGenerator(batch_size=remainder_shape,
                                                           shuffle=False,
                                                           steps_per_epoch=1,
                                                           data=[norm_data_remainder], 
                                                           pbar=pbar)
                predictions[data_gen_shape:] = np.asarray(self.keras_model.predict(remainder_generator))

        if self.labels_normalizer is not None:
            predictions = self.labels_normalizer.denormalize(list_to_dict(self.keras_model.output_names, predictions))
        else:
            predictions *= self.labels_std
            predictions += self.labels_mean

        return predictions['output']

    def evaluate(self, input_data, labels):
        """
        Evaluate neural network by provided input data and labels and get back a metrics score

        :param input_data: Data to be inferred with neural network
        :type input_data: ndarray
        :param labels: labels
        :type labels: ndarray
        :return: metrics score dictionary
        :rtype: dict
        :History: 2018-May-20 - Written - Henry Leung (University of Toronto)
        """
        self.has_model_check()
        input_data = list_to_dict(self.keras_model.input_names, input_data)
        labels = list_to_dict(self.keras_model.output_names, labels)

        # check if exists (existing means the model has already been trained (e.g. fine-tuning), so we do not need calculate mean/std again)
        if self.input_normalizer is None:
            self.input_normalizer = Normalizer(mode=self.input_norm_mode)
            self.labels_normalizer = Normalizer(mode=self.labels_norm_mode)

            norm_data = self.input_normalizer.normalize(input_data)
            self.input_mean, self.input_std = self.input_normalizer.mean_labels, self.input_normalizer.std_labels
            norm_labels = self.labels_normalizer.normalize(labels)
            self.labels_mean, self.labels_std = self.labels_normalizer.mean_labels, self.labels_normalizer.std_labels
        else:
            norm_data = self.input_normalizer.normalize(input_data, calc=False)
            norm_labels = self.labels_normalizer.normalize(labels, calc=False)

        norm_data = self._tensor_dict_sanitize(norm_data, self.keras_model.input_names)
        norm_labels = self._tensor_dict_sanitize(norm_labels, self.keras_model.output_names)

        total_num = input_data['input'].shape[0]
        eval_batchsize = self.batch_size if total_num > self.batch_size else total_num
        steps = total_num // self.batch_size if total_num > self.batch_size else 1

        start_time = time.time()
        print("Starting Evaluation")

        evaluate_generator = CNNDataGenerator(batch_size=eval_batchsize,
                                              shuffle=False,
                                              steps_per_epoch=steps,
                                              data=[norm_data, norm_labels])

        scores = self.keras_model.evaluate(evaluate_generator)
        if isinstance(scores, float):  # make sure scores is iterable
            scores = list(str(scores))
        outputname = self.keras_model.output_names
        funcname = self.keras_model.metrics_names

        print(f'Completed Evaluation, {(time.time() - start_time):.{2}f}s elapsed')

        return list_to_dict(funcname, scores)

    @deprecated_copy_signature(fit)
    def train(self, *args, **kwargs):
        return self.fit(*args, **kwargs)

    @deprecated_copy_signature(fit_on_batch)
    def train_on_batch(self, *args, **kwargs):
        return self.fit_on_batch(*args, **kwargs)
    
    @deprecated_copy_signature(predict)
    def test(self, *args, **kwargs):
        return self.predict(*args, **kwargs) class ConvVAEBase(NeuralNetMaster, ABC):
    """
    Top-level class for a Convolutional Variational Autoencoder

    :History: 2018-Jan-06 - Written - Henry Leung (University of Toronto)
    """

    def __init__(self):
        super().__init__()
        self.name = 'Convolutional Variational Autoencoder'
        self._model_type = 'CVAE'
        self.initializer = None
        self.activation = None
        self._last_layer_activation = None
        self.num_filters = None
        self.filter_len = None
        self.pool_length = None
        self.num_hidden = None
        self.reduce_lr_epsilon = None
        self.reduce_lr_min = None
        self.reduce_lr_patience = None
        self.l1 = None
        self.l2 = None
        self.maxnorm = None
        self.latent_dim = None
        self.val_size = 0.1
        self.dropout_rate = 0.0

        self.keras_vae = None
        self.keras_encoder = None
        self.keras_decoder = None
        self.loss = None

        self._input_shape = None

        self.input_norm_mode = 255
        self.labels_norm_mode = 255
        self.input_mean = None
        self.input_std = None
        self.labels_mean = None
        self.labels_std = None

    def compile(self,
                optimizer=None,
                loss=None,
                metrics=None,
                weighted_metrics=None,
                loss_weights=None,
                sample_weight_mode=None):
        self.keras_model, self.keras_encoder, self.keras_decoder = self.model()

        if optimizer is not None:
            self.optimizer = optimizer
        elif self.optimizer is None or self.optimizer == 'adam':
            self.optimizer = Adam(learning_rate=self.lr, beta_1=self.beta_1, beta_2=self.beta_2, epsilon=self.optimizer_epsilon,
                                  decay=0.0)
        if metrics is not None:
            self.metrics = metrics
        self.loss = mean_squared_error if not (loss and self.loss) else loss
        self.metrics = [mean_absolute_error, mean_error] if not self.metrics else self.metrics

        self.keras_model.compile(loss=self.loss,
                                 optimizer=self.optimizer,
                                 metrics=self.metrics,
                                 weighted_metrics=weighted_metrics,
                                 loss_weights=loss_weights,
                                 sample_weight_mode=sample_weight_mode)

        # inject custom training step if needed
        try:
            self.custom_train_step()
        except NotImplementedError:
            pass
        except TypeError:
            self.keras_model.train_step = self.custom_train_step
       # inject custom testing  step if needed
        try:
            self.custom_test_step()
        except NotImplementedError:
            pass
        except TypeError:
            self.keras_model.test_step = self.custom_test_step

        return None

    def pre_training_checklist_child(self, input_data, input_recon_target, sample_weights):
        if self.task == 'classification':
            raise RuntimeError('astroNN VAE does not support classification task')
        elif self.task == 'binary_classification':
            raise RuntimeError('astroNN VAE does not support binary classification task')

        input_data, input_recon_target = self.pre_training_checklist_master(input_data, input_recon_target)

        if isinstance(input_data, H5Loader):
            self.targetname = input_data.target
            input_data, input_recon_target = input_data.load()

        # check if exists (existing means the model has already been trained (e.g. fine-tuning), so we do not need calculate mean/std again)
        if self.input_normalizer is None:
            self.input_normalizer = Normalizer(mode=self.input_norm_mode)
            self.labels_normalizer = Normalizer(mode=self.labels_norm_mode)

            norm_data = self.input_normalizer.normalize(input_data)
            self.input_mean, self.input_std = self.input_normalizer.mean_labels, self.input_normalizer.std_labels
            norm_labels = self.labels_normalizer.normalize(input_recon_target)
            self.labels_mean, self.labels_std = self.labels_normalizer.mean_labels, self.labels_normalizer.std_labels
        else:
            norm_data = self.input_normalizer.normalize(input_data, calc=False)
            norm_labels = self.labels_normalizer.normalize(input_recon_target, calc=False)

        if self.keras_model is None:  # only compile if there is no keras_model, e.g. fine-tuning does not required
            self.compile()
        
        norm_data = self._tensor_dict_sanitize(norm_data, self.keras_model.input_names)
        norm_labels = self._tensor_dict_sanitize(norm_labels, self.keras_model.output_names)

        self.train_idx, self.val_idx = train_test_split(np.arange(self.num_train + self.val_num),
                                                        test_size=self.val_size)

        norm_data_training = {}
        norm_data_val = {}
        norm_labels_training = {}
        norm_labels_val = {}
        for name in norm_data.keys():
            norm_data_training.update({name: norm_data[name][self.train_idx]})
            norm_data_val.update({name: norm_data[name][self.val_idx]})
        for name in norm_labels.keys():
            norm_labels_training.update({name: norm_labels[name][self.train_idx]})
            norm_labels_val.update({name: norm_labels[name][self.val_idx]})

        if sample_weights is not None:        
            sample_weights_training = sample_weights[self.train_idx]
            sample_weights_val = sample_weights[self.val_idx]
        else:
            sample_weights_training = None
            sample_weights_val = None

        self.training_generator = CVAEDataGenerator(batch_size=self.batch_size,
                                                    shuffle=True,
                                                    steps_per_epoch=self.num_train // self.batch_size,
                                                    data=[norm_data_training,
                                                          norm_labels_training],
                                                    manual_reset=False, 
                                                    sample_weights=sample_weights_training)

        val_batchsize = self.batch_size if len(self.val_idx) > self.batch_size else len(self.val_idx)
        self.validation_generator = CVAEDataGenerator(batch_size=val_batchsize,
                                                      shuffle=True,
                                                      steps_per_epoch=max(self.val_num // self.batch_size, 1),
                                                      data=[norm_data_val,
                                                            norm_labels_val],
                                                      manual_reset=True, 
                                                      sample_weights=sample_weights_val)

        return input_data, input_recon_target

    def fit(self, input_data, input_recon_target, sample_weights=None):
        """
        Train a Convolutional Autoencoder

        :param input_data: Data to be trained with neural network
        :type input_data: ndarray
        :param input_recon_target: Data to be reconstructed
        :type input_recon_target: ndarray
        :param sample_weights: Sample weights (if any)
        :type sample_weights: Union([NoneType, ndarray])
        :return: None
        :rtype: NoneType
        :History: 2017-Dec-06 - Written - Henry Leung (University of Toronto)
        """

        # Call the checklist to create astroNN folder and save parameters
        self.pre_training_checklist_child(input_data, input_recon_target, sample_weights)

        reduce_lr = ReduceLROnPlateau(monitor='val_output_loss', factor=0.5, min_delta=self.reduce_lr_epsilon,
                                      patience=self.reduce_lr_patience, min_lr=self.reduce_lr_min, mode='min',
                                      verbose=2)

        self.virtual_cvslogger = VirutalCSVLogger()

        self.__callbacks = [reduce_lr, self.virtual_cvslogger]  # default must have unchangeable callbacks

        if self.callbacks is not None:
            if isinstance(self.callbacks, list):
                self.__callbacks.extend(self.callbacks)
            else:
                self.__callbacks.append(self.callbacks)

        start_time = time.time()

        self.keras_model.fit(self.training_generator,
                             validation_data=self.validation_generator,
                             epochs=self.max_epochs, verbose=self.verbose, workers=os.cpu_count(),
                             callbacks=self.__callbacks,
                             use_multiprocessing=MULTIPROCESS_FLAG)

        print(f'Completed Training, {(time.time() - start_time):.{2}f}s in total')

        if self.autosave is True:
            # Call the post training checklist to save parameters
            self.save()

        return None

    def fit_on_batch(self, input_data, input_recon_target, sample_weights=None):
        """
        Train a AutoEncoder by running a single gradient update on all of your data, suitable for fine-tuning

        :param input_data: Data to be trained with neural network
        :type input_data: ndarray
        :param input_recon_target: Data to be reconstructed
        :type input_recon_target: ndarray
        :param sample_weights: Sample weights (if any)
        :type sample_weights: Union([NoneType, ndarray])
        :return: None
        :rtype: NoneType
        :History: 2018-Aug-25 - Written - Henry Leung (University of Toronto)
        """

        input_data, input_recon_target = self.pre_training_checklist_master(input_data, input_recon_target)

        # check if exists (existing means the model has already been trained (e.g. fine-tuning), so we do not need calculate mean/std again)
        if self.input_normalizer is None:
            self.input_normalizer = Normalizer(mode=self.input_norm_mode)
            self.labels_normalizer = Normalizer(mode=self.labels_norm_mode)

            norm_data = self.input_normalizer.normalize(input_data)
            self.input_mean, self.input_std = self.input_normalizer.mean_labels, self.input_normalizer.std_labels
            norm_labels = self.labels_normalizer.normalize(input_recon_target)
            self.labels_mean, self.labels_std = self.labels_normalizer.mean_labels, self.labels_normalizer.std_labels
        else:
            norm_data = self.input_normalizer.normalize(input_data, calc=False)
            norm_labels = self.labels_normalizer.normalize(input_recon_target, calc=False)
            
        norm_data = self._tensor_dict_sanitize(norm_data, self.keras_model.input_names)
        norm_labels = self._tensor_dict_sanitize(norm_labels, self.keras_model.output_names)

        start_time = time.time()

        fit_generator = CVAEDataGenerator(batch_size=input_data['input'].shape[0],
                                          shuffle=False,
                                          steps_per_epoch=1,
                                          data=[norm_data,
                                                norm_labels], 
                                          sample_weights=sample_weights)

        scores = self.keras_model.fit(fit_generator,
                                      epochs=1,
                                      verbose=self.verbose,
                                      workers=os.cpu_count(),
                                      use_multiprocessing=MULTIPROCESS_FLAG)

        print(f'Completed Training on Batch, {(time.time() - start_time):.{2}f}s in total')

        return None

    def post_training_checklist_child(self):
        self.keras_model.save(self.fullfilepath + _astroNN_MODEL_NAME)
        print(_astroNN_MODEL_NAME + f' saved to {(self.fullfilepath + _astroNN_MODEL_NAME)}')

        self.hyper_txt.write(f"Dropout Rate: {self.dropout_rate} \n")
        self.hyper_txt.flush()
        self.hyper_txt.close()

        data = {'id': self.__class__.__name__,
                'pool_length': self.pool_length,
                'filterlen': self.filter_len,
                'filternum': self.num_filters,
                'hidden': self.num_hidden,
                'input': self._input_shape,
                'labels': self._labels_shape,
                'task': self.task,
                'activation': self.activation,
                'input_mean': dict_np_to_dict_list(self.input_mean),
                'labels_mean': dict_np_to_dict_list(self.labels_mean),
                'input_std': dict_np_to_dict_list(self.input_std),
                'labels_std': dict_np_to_dict_list(self.labels_std),
                'valsize': self.val_size,
                'targetname': self.targetname,
                'dropout_rate': self.dropout_rate,
                'l1': self.l1, 'l2': self.l2,
                'maxnorm': self.maxnorm,
                'input_norm_mode': self.input_normalizer.normalization_mode,
                'labels_norm_mode': self.labels_normalizer.normalization_mode,
                'batch_size': self.batch_size,
                'latent': self.latent_dim}

        with open(self.fullfilepath + '/astroNN_model_parameter.json', 'w') as f:
            json.dump(data, f, indent=4, sort_keys=True)

    def predict(self, input_data):
        """
        Use the neural network to do inference and get reconstructed data

        :param input_data: Data to be inferred with neural network
        :type input_data: ndarray
        :return: reconstructed data
        :rtype: ndarry
        :History: 2017-Dec-06 - Written - Henry Leung (University of Toronto)
        """
        input_data = self.pre_testing_checklist_master(input_data)

        if self.input_normalizer is not None:
            input_array = self.input_normalizer.normalize(input_data, calc=False)
        else:
            # Prevent shallow copy issue
            input_array = np.array(input_data)
            input_array -= self.input_mean
            input_array /= self.input_std

        total_test_num = input_data['input'].shape[0]  # Number of testing data

        # for number of training data smaller than batch_size
        if total_test_num < self.batch_size:
            self.batch_size = total_test_num

        # Due to the nature of how generator works, no overlapped prediction
        data_gen_shape = (total_test_num // self.batch_size) * self.batch_size
        remainder_shape = total_test_num - data_gen_shape  # Remainder from generator

        predictions = np.zeros((total_test_num, self._labels_shape['output'], 1))

        norm_data_main = {}
        norm_data_remainder = {}
        for name in input_array.keys():
            norm_data_main.update({name: input_array[name][:data_gen_shape]})
            norm_data_remainder.update({name: input_array[name][data_gen_shape:]})
            
        norm_data_main = self._tensor_dict_sanitize(norm_data_main, self.keras_model.input_names)
        norm_data_remainder = self._tensor_dict_sanitize(norm_data_remainder, self.keras_model.input_names)

        start_time = time.time()
        print("Starting Inference")

        # Data Generator for prediction
        with tqdm(total=total_test_num, unit="sample") as pbar:
            prediction_generator = CVAEPredDataGenerator(batch_size=self.batch_size,
                                                        shuffle=False,
                                                        steps_per_epoch=total_test_num // self.batch_size,
                                                        data=[norm_data_main])
            result = np.asarray(self.keras_model.predict(prediction_generator))

            if remainder_shape != 0:
                remainder_generator = CVAEPredDataGenerator(batch_size=remainder_shape,
                                                            shuffle=False,
                                                            steps_per_epoch=1,
                                                            data=[norm_data_remainder], 
                                                            pbar=pbar)
                remainder_result = np.asarray(self.keras_model.predict(remainder_generator))
                result = np.concatenate((result, remainder_result))

        if self.labels_normalizer is not None:
            # TODO: handle named output in the future
            predictions[:, :, 0] = self.labels_normalizer.denormalize(list_to_dict(self.keras_model.output_names,
                                                                                   predictions[:, :, 0]))['output']
        else:
            predictions[:, :, 0] *= self.labels_std
            predictions[:, :, 0] += self.labels_mean

        print(f'Completed Inference, {(time.time() - start_time):.{2}f}s elapsed')

        return predictions

    def predict_encoder(self, input_data):
        """
        Use the neural network to do inference and get the hidden layer encoding/representation

        :param input_data: Data to be inferred with neural network
        :type input_data: ndarray
        :return: hidden layer encoding/representation
        :rtype: ndarray
        :History: 2017-Dec-06 - Written - Henry Leung (University of Toronto)
        """
        input_data = self.pre_testing_checklist_master(input_data)
        # Prevent shallow copy issue
        if self.input_normalizer is not None:
            input_array = self.input_normalizer.normalize(input_data, calc=False)
        else:
            # Prevent shallow copy issue
            input_array = np.array(input_data)
            input_array -= self.input_mean
            input_array /= self.input_std

        total_test_num = input_data['input'].shape[0]  # Number of testing data

        # for number of training data smaller than batch_size
        if total_test_num < self.batch_size:
            self.batch_size = input_data.shape[0]

        # Due to the nature of how generator works, no overlapped prediction
        data_gen_shape = (total_test_num // self.batch_size) * self.batch_size
        remainder_shape = total_test_num - data_gen_shape  # Remainder from generator

        norm_data_main = {}
        norm_data_remainder = {}
        for name in input_array.keys():
            norm_data_main.update({name: input_array[name][:data_gen_shape]})
            norm_data_remainder.update({name: input_array[name][data_gen_shape:]})

        encoding = np.zeros((total_test_num, self.latent_dim))

        start_time = time.time()
        print("Starting Inference on Encoder")

        # Data Generator for prediction
        prediction_generator = CVAEPredDataGenerator(batch_size=self.batch_size,
                                                     shuffle=False,
                                                     steps_per_epoch=total_test_num // self.batch_size,
                                                     data=[norm_data_main])
        encoding[:data_gen_shape] = np.asarray(self.keras_encoder.predict(prediction_generator))

        if remainder_shape != 0:
            # assume its caused by mono images, so need to expand dim by 1
            for name in input_array.keys():
                if len(norm_data_remainder[name][0].shape) != len(self._input_shape[name]):
                    norm_data_remainder.update({name: np.expand_dims(norm_data_remainder[name], axis=-1)})
            result = self.keras_encoder.predict(norm_data_remainder)
            encoding[data_gen_shape:] = result

        print(f'Completed Inference on Encoder, {(time.time() - start_time):.{2}f}s elapsed')

        return encoding

    def evaluate(self, input_data, labels):
        """
        Evaluate neural network by provided input data and labels/reconstruction target to get back a metrics score

        :param input_data: Data to be inferred with neural network
        :type input_data: ndarray
        :param labels: labels
        :type labels: ndarray
        :return: metrics score
        :rtype: float
        :History: 2018-May-20 - Written - Henry Leung (University of Toronto)
        """
        self.has_model_check()
        input_data = list_to_dict(self.keras_model.input_names, input_data)
        labels = list_to_dict(self.keras_model.output_names, labels)

        # check if exists (existing means the model has already been trained (e.g. fine-tuning), so we do not need calculate mean/std again)
        if self.input_normalizer is None:
            self.input_normalizer = Normalizer(mode=self.input_norm_mode)
            self.labels_normalizer = Normalizer(mode=self.labels_norm_mode)

            norm_data = self.input_normalizer.normalize(input_data)
            self.input_mean, self.input_std = self.input_normalizer.mean_labels, self.input_normalizer.std_labels
            norm_labels = self.labels_normalizer.normalize(labels)
            self.labels_mean, self.labels_std = self.labels_normalizer.mean_labels, self.labels_normalizer.std_labels
        else:
            norm_data = self.input_normalizer.normalize(input_data, calc=False)
            norm_labels = self.labels_normalizer.normalize(labels, calc=False)
            
        norm_data = self._tensor_dict_sanitize(norm_data, self.keras_model.input_names)
        norm_labels = self._tensor_dict_sanitize(norm_labels, self.keras_model.input_names)

        total_num = input_data['input'].shape[0]
        eval_batchsize = self.batch_size if total_num > self.batch_size else total_num
        steps = total_num // self.batch_size if total_num > self.batch_size else 1

        start_time = time.time()
        print("Starting Evaluation")

        evaluate_generator = CVAEDataGenerator(batch_size=eval_batchsize,
                                               shuffle=False,
                                               steps_per_epoch=steps,
                                               data=[norm_data,
                                                     norm_labels])

        scores = self.keras_model.evaluate(evaluate_generator)
        if isinstance(scores, float):  # make sure scores is iterable
            scores = list(str(scores))
        outputname = self.keras_model.output_names
        funcname = self.keras_model.metrics_names

        print(f'Completed Evaluation, {(time.time() - start_time):.{2}f}s elapsed')

        return list_to_dict(funcname, scores)

    @deprecated_copy_signature(fit)
    def train(self, *args, **kwargs):
        return self.fit(*args, **kwargs)

    @deprecated_copy_signature(fit_on_batch)
    def train_on_batch(self, *args, **kwargs):
        return self.fit_on_batch(*args, **kwargs)
    
    @deprecated_copy_signature(predict)
    def test(self, *args, **kwargs):
        return self.predict(*args, **kwargs)
    
    @deprecated_copy_signature(predict)
    def test_encoder(self, *args, **kwargs):
        return self.predict_encoder(*args, **kwargs) class MCDropout(Layer):
    """
    Dropout Layer for Bayesian Neural Network, this layer will always on regardless the learning phase flag

    :param rate: Dropout Rate between 0 and 1
    :type rate: float
    :param disable: Dropout on or off
    :type disable: boolean
    :return: A layer
    :rtype: object
    :History: 2018-Feb-05 - Written - Henry Leung (University of Toronto)
    """

    def __init__(self, rate, disable=False, noise_shape=None, name=None, **kwargs):
        self.rate = rate
        self.disable_layer = disable
        self.supports_masking = True
        self.noise_shape = noise_shape
        if not name:
            prefix = self.__class__.__name__
            name = prefix + '_' + str(tfk.backend.get_uid(prefix))
        super().__init__(name=name, **kwargs)

    def _get_noise_shape(self, inputs):
        if self.noise_shape is None:
            return self.noise_shape

        symbolic_shape = tf.shape(inputs)
        noise_shape = [symbolic_shape[axis] if shape is None else shape
                       for axis, shape in enumerate(self.noise_shape)]
        return tuple(noise_shape)

    def call(self, inputs, training=None):
        """
        :Note: Equivalent to __call__()
        :param inputs: Tensor to be applied
        :type inputs: tf.Tensor
        :return: Tensor after applying the layer
        :rtype: tf.Tensor
        """
        noise_shape = self._get_noise_shape(inputs)
        if self.disable_layer is True:
            return inputs
        else:
            return tf.nn.dropout(x=inputs,
                                 rate=self.rate,
                                 noise_shape=noise_shape)

    def get_config(self):
        """
        :return: Dictionary of configuration
        :rtype: dict
        """
        config = {'rate': self.rate,
                  'noise_shape': self.noise_shape}
        base_config = super().get_config()
        return {**dict(base_config.items()), **config}

    def compute_output_shape(self, input_shape):
        return input_shape class BoolMask(Layer):
    """
    Boolean Masking layer, please notice it is best to flatten input before using BoolMask

    :param mask: numpy boolean array as a mask for incoming tensor
    :type mask: np.ndarray
    :return: A layer
    :rtype: object
    :History: 2018-May-28 - Written - Henry Leung (University of Toronto)
    """

    def __init__(self, mask, name=None, **kwargs):
        if sum(mask) == 0:
            raise ValueError("The mask is all False, which is invalid")
        else:
            self.boolmask = mask
        self.mask_shape = self.boolmask.sum()
        self.supports_masking = True
        if not name:
            prefix = self.__class__.__name__
            name = prefix + '_' + str(tfk.backend.get_uid(prefix))
        super().__init__(name=name, **kwargs)

    def compute_output_shape(self, input_shape):
        input_shape = tensor_shape.TensorShape(input_shape)
        input_shape = input_shape.with_rank_at_least(2)
        return input_shape[:-1].concatenate(self.mask_shape)

    def call(self, inputs, training=None):
        """
        :Note: Equivalent to __call__()
        :param inputs: Tensor to be applied
        :type inputs: tf.Tensor
        :return: Tensor after applying the layer which is just the masked tensor
        :rtype: tf.Tensor
        """
        batchsize = tf.shape(inputs)[0]
        # need to reshape because tf.keras cannot get the Tensor shape correctly from tf.boolean_mask op
        return tf.reshape(tf.boolean_mask(inputs, self.boolmask, axis=1), [batchsize, self.mask_shape])

    def get_config(self):
        """
        :return: Dictionary of configuration
        :rtype: dict
        """
        config = {'None': None}
        base_config = super().get_config()
        return {**dict(base_config.items()), **config} class StopGrad(Layer):
    """
    Stop gradient backpropagation via this layer during training, act as an identity layer during testing by default.

    :param always_on: Default False which will on during train time and off during test time. True to enable it in every situation
    :type always_on: bool
    :return: A layer
    :rtype: object
    :History: 2018-May-23 - Written - Henry Leung (University of Toronto)
    """

    def __init__(self, name=None, always_on=False, **kwargs):
        if not name:
            prefix = self.__class__.__name__
            name = prefix + '_' + str(tfk.backend.get_uid(prefix))
        super().__init__(name=name, **kwargs)
        self.always_on = always_on

    def compute_output_shape(self, input_shape):
        return input_shape

    def call(self, inputs, training=None):
        """
        :Note: Equivalent to __call__()
        :param inputs: Tensor to be applied
        :type inputs: tf.Tensor
        :return: Tensor after applying the layer which is just the original tensor
        :rtype: tf.Tensor
        """
        if self.always_on:
            return tf.stop_gradient(inputs)
        else:
            if training is None:
                training = tfk.backend.learning_phase()
            output_tensor = tf.where(tf.equal(training, True), tf.stop_gradient(inputs), inputs)
            output_tensor._uses_learning_phase = True
            return output_tensor

    def get_config(self):
        """
        :return: Dictionary of configuration
        :rtype: dict
        """
        config = {'None': None}
        base_config = super().get_config()
        return {**dict(base_config.items()), **config} class KLDivergenceLayer(Layer):
    """
    | Identity transform layer that adds KL divergence to the final model losses.
    | KL divergence used to force the latent space match the prior (in this case its unit gaussian)

    :return: A layer
    :rtype: object
    :History: 2018-Feb-05 - Written - Henry Leung (University of Toronto)
    """

    def __init__(self, name=None, **kwargs):
        self.is_placeholder = True
        if not name:
            prefix = self.__class__.__name__
            name = prefix + '_' + str(tfk.backend.get_uid(prefix))
        super().__init__(name=name, **kwargs)

    def call(self, inputs, training=None):
        """
        :Note: Equivalent to __call__()
        :param inputs: Tensor to be applied, concatenated tf.tensor of mean and std in latent space
        :type inputs: tf.Tensor
        :return: Tensor after applying the layer
        :rtype: tf.Tensor
        """
        mu, log_var = inputs
        kl_batch = - .5 * tf.reduce_sum(1 + log_var - tf.square(mu) - tf.exp(log_var), axis=-1)
        self.add_loss(tf.reduce_mean(kl_batch), inputs=inputs)

        return inputs

    def get_config(self):
        """
        :return: Dictionary of configuration
        :rtype: dict
        """
        config = {'None': None}
        base_config = super().get_config()
        return {**dict(base_config.items()), **config}

    def compute_output_shape(self, input_shape):
        return input_shape class TensorInput(Layer):
    """
    TensorInput layer

    :param tensor: tensor, usually is a tensor generating random number
    :type tensor: tf.Tensor
    :return: A layer
    :rtype: object
    :History: 2020-May-3 - Written - Henry Leung (University of Toronto)
    """

    def __init__(self, tensor=None, name=None, **kwargs):
        self.supports_masking = True
        self.tensor = tensor
        if not name:
            prefix = self.__class__.__name__
            name = prefix + '_' + str(tfk.backend.get_uid(prefix))
        super().__init__(name=name, **kwargs)

    def compute_output_shape(self, input_shape):
        return self.input_shape

    def call(self, inputs, training=None):
        """
        :Note: Equivalent to __call__()
        :param inputs: Tensor to be applied
        :type inputs: tf.Tensor
        :return: Tensor after applying the layer which is just the masked tensor
        :rtype: tf.Tensor
        """
        return self.tensor

    def get_config(self):
        """
        :return: Dictionary of configuration
        :rtype: dict
        """
        config = {'None': None}
        base_config = super().get_config()
        return {**dict(base_config.items()), **config} def bayesian_binary_crossentropy_wrapper(logit_var):
    """
    | Binary crossentropy between an output tensor and a target tensor for Bayesian Neural Network
    | equation (12) of arxiv:1703.04977

    :param logit_var: Predictive variance
    :type logit_var: Union(tf.Tensor, tf.Variable)
    :return: Robust binary_crossentropy function for predictive variance neurones which matches Keras losses API
    :rtype: function
    :Returned Function Parameter:
            | **function(y_true, y_pred)**
            |   - **y_true** (*tf.Tensor*): Ground Truth
            |   - **y_pred** (*tf.Tensor*): Prediction in logits space
            |   Return (*tf.Tensor*): Robust binary crossentropy
    :History: 2018-Mar-15 - Written - Henry Leung (University of Toronto)
    """

    # y_pred is logits
    def bayesian_crossentropy(y_true, y_pred, sample_weight=None):
        return robust_binary_crossentropy(y_true, y_pred, logit_var, sample_weight)

    # set the name to be the same as parent so it can be found
    bayesian_crossentropy.__name__ = 'bayesian_binary_crossentropy_wrapper'

    return bayesian_crossentropy def bayesian_binary_crossentropy_var_wrapper(logits):
    """
    | Binary crossentropy between an output tensor and a target tensor for Bayesian Neural Network
    | equation (12) of arxiv:1703.04977

    :param logits: Prediction in logits space
    :type logits: Union(tf.Tensor, tf.Variable)
    :return: Robust binary_crossentropy function for predictive variance neurones which matches Keras losses API
    :rtype: function
    :Returned Function Parameter:
            | **function(y_true, y_pred)**
            |   - **y_true** (*tf.Tensor*): Ground Truth
            |   - **y_pred** (*tf.Tensor*): Predictive variance in logits space
            |   Return (*tf.Tensor*): Robust binary crossentropy
    :History: 2018-Mar-15 - Written - Henry Leung (University of Toronto)
    """

    # y_pred is predictive entropy
    def bayesian_crossentropy(y_true, y_pred, sample_weight=None):
        return robust_binary_crossentropy(y_true, logits, y_pred, sample_weight)

    # set the name to be the same as parent so it can be found
    bayesian_crossentropy.__name__ = 'bayesian_binary_crossentropy_var_wrapper'

    return bayesian_crossentropy def bayesian_categorical_crossentropy_wrapper(logit_var):
    """
    | Categorical crossentropy between an output tensor and a target tensor for Bayesian Neural Network
    | equation (12) of arxiv:1703.04977

    :param logit_var: Predictive variance
    :type logit_var: Union(tf.Tensor, tf.Variable)
    :return: Robust categorical_crossentropy function for predictive variance neurones which matches Keras losses API
    :rtype: function
    :Returned Function Parameter:
            | **function(y_true, y_pred)**
            |   - **y_true** (*tf.Tensor*): Ground Truth
            |   - **y_pred** (*tf.Tensor*): Prediction in logits space
            |   Return (*tf.Tensor*): Robust categorical crossentropy
    :History: 2018-Mar-15 - Written - Henry Leung (University of Toronto)
    """

    # y_pred is logits
    def bayesian_crossentropy(y_true, y_pred, sample_weight=None):
        return robust_categorical_crossentropy(y_true, y_pred, logit_var, sample_weight)

    # set the name to be the same as parent so it can be found
    bayesian_crossentropy.__name__ = 'bayesian_categorical_crossentropy_wrapper'

    return bayesian_crossentropy def bayesian_categorical_crossentropy_var_wrapper(logits):
    """
    | Categorical crossentropy between an output tensor and a target tensor for Bayesian Neural Network
    | equation (12) of arxiv:1703.04977

    :param logits: Prediction in logits space
    :type logits: Union(tf.Tensor, tf.Variable)
    :return: Robust categorical_crossentropy function for predictive variance neurones which matches Keras losses API
    :rtype: function
    :Returned Function Parameter:
            | **function(y_true, y_pred)**
            |   - **y_true** (*tf.Tensor*): Ground Truth
            |   - **y_pred** (*tf.Tensor*): Predictive variance in logits space
            |   Return (*tf.Tensor*): Robust categorical crossentropy
    :History: 2018-Mar-15 - Written - Henry Leung (University of Toronto)
    """

    # y_pred is predictive entropy
    def bayesian_crossentropy(y_true, y_pred, sample_weight=None):
        return robust_categorical_crossentropy(y_true, logits, y_pred, sample_weight)

    # set the name to be the same as parent so it can be found
    bayesian_crossentropy.__name__ = 'bayesian_categorical_crossentropy_var_wrapper'

    return bayesian_crossentropy def mse_lin_wrapper(var, labels_err):
    """
    Calculate predictive variance, and takes account of labels error in Bayesian Neural Network

    :param var: Predictive Variance
    :type var: Union(tf.Tensor, tf.Variable)
    :param labels_err: Known labels error, give zeros if unknown/unavailable
    :type labels_err: Union(tf.Tensor, tf.Variable)
    :return: Robust MSE function for labels prediction neurones, which matches Keras losses API
    :rtype: function
    :Returned Funtion Parameter:
            | **function(y_true, y_pred)**
            |   - **y_true** (*tf.Tensor*): Ground Truth
            |   - **y_pred** (*tf.Tensor*): Prediction
            |   Return (*tf.Tensor*): Robust Mean Squared Error
    :History: 2017-Nov-16 - Written - Henry Leung (University of Toronto)
    """

    def mse_lin(y_true, y_pred, sample_weight=None):
        return robust_mse(y_true, y_pred, var, labels_err, sample_weight)

    mse_lin.__name__ = 'mse_lin_wrapper'  # set the name to be the same as parent so it can be found

    return mse_lin def mse_var_wrapper(lin, labels_err):
    """
    Calculate predictive variance, and takes account of labels error in Bayesian Neural Network

    :param lin: Prediction
    :type lin: Union(tf.Tensor, tf.Variable)
    :param labels_err: Known labels error, give zeros if unknown/unavailable
    :type labels_err: Union(tf.Tensor, tf.Variable)
    :return: Robust MSE function for predictive variance neurones which matches Keras losses API
    :rtype: function
    :Returned Funtion Parameter:
            | **function(y_true, y_pred)**
            |   - **y_true** (*tf.Tensor*): Ground Truth
            |   - **y_pred** (*tf.Tensor*): Predictive Variance
            |   Return (*tf.Tensor*): Robust Mean Squared Error
    :History: 2017-Nov-16 - Written - Henry Leung (University of Toronto)
    """

    def mse_var(y_true, y_pred, sample_weight=None):
        return robust_mse(y_true, lin, y_pred, labels_err, sample_weight)

    mse_var.__name__ = 'mse_var_wrapper'  # set the name to be the same as parent so it can be found

    return mse_var
import numpy as np
import tensorflow as tf
from tensorflow import keras as tfk
from astroNN.apogee import aspcap_mask
from astroNN.models.base_bayesian_cnn import BayesianCNNBase
from astroNN.models.base_cnn import CNNBase
from astroNN.models.base_vae import ConvVAEBase
from astroNN.nn.layers import MCDropout, BoolMask, StopGrad, KLDivergenceLayer, TensorInput
from astroNN.nn.losses import bayesian_binary_crossentropy_wrapper, bayesian_binary_crossentropy_var_wrapper
from astroNN.nn.losses import bayesian_categorical_crossentropy_wrapper, bayesian_categorical_crossentropy_var_wrapper
from astroNN.nn.losses import mse_lin_wrapper, mse_var_wrapper
    """
    Class for Bayesian convolutional neural network for stellar spectra analysis

    :History: 2017-Dec-21 - Written - Henry Leung (University of Toronto)
    """

    def __init__(self, lr=0.0005, dropout_rate=0.3):
        super().__init__()

        self._implementation_version = '1.0'
        self.initializer = RandomNormal(mean=0.0, stddev=0.05)
        self.activation = 'relu'
        self.num_filters = [2, 4]
        self.filter_len = 8
        self.pool_length = 4
        self.num_hidden = [196, 96]
        self.max_epochs = 100
        self.lr = lr
        self.reduce_lr_epsilon = 0.00005

        self.reduce_lr_min = 1e-8
        self.reduce_lr_patience = 2
        self.l2 = 5e-9
        self.dropout_rate = dropout_rate

        self.input_norm_mode = 3

        self.task = 'regression'

        self.targetname = ['teff', 'logg', 'M', 'alpha', 'C', 'C1', 'N', 'O', 'Na', 'Mg', 'Al', 'Si', 'P', 'S', 'K',
                           'Ca', 'Ti', 'Ti2', 'V', 'Cr', 'Mn', 'Fe', 'Co', 'Ni', 'fakemag']

    def model(self):
        input_tensor = Input(shape=self._input_shape['input'], name='input')
        labels_err_tensor = Input(shape=(self._labels_shape['output'],), name='labels_err')

        cnn_layer_1 = Conv1D(kernel_initializer=self.initializer, padding="same", filters=self.num_filters[0],
                             kernel_size=self.filter_len, kernel_regularizer=regularizers.l2(self.l2))(input_tensor)
        activation_1 = Activation(activation=self.activation)(cnn_layer_1)
        dropout_1 = MCDropout(self.dropout_rate, disable=self.disable_dropout)(activation_1)
        cnn_layer_2 = Conv1D(kernel_initializer=self.initializer, padding="same", filters=self.num_filters[1],
                             kernel_size=self.filter_len, kernel_regularizer=regularizers.l2(self.l2))(dropout_1)
        activation_2 = Activation(activation=self.activation)(cnn_layer_2)
        maxpool_1 = MaxPooling1D(pool_size=self.pool_length)(activation_2)
        flattener = Flatten()(maxpool_1)
        dropout_2 = MCDropout(self.dropout_rate, disable=self.disable_dropout)(flattener)
        layer_3 = Dense(units=self.num_hidden[0], kernel_regularizer=regularizers.l2(self.l2),
                        kernel_initializer=self.initializer)(dropout_2)
        activation_3 = Activation(activation=self.activation)(layer_3)
        dropout_3 = MCDropout(self.dropout_rate, disable=self.disable_dropout)(activation_3)
        layer_4 = Dense(units=self.num_hidden[1], kernel_regularizer=regularizers.l2(self.l2),
                        kernel_initializer=self.initializer)(dropout_3)
        activation_4 = Activation(activation=self.activation)(layer_4)
        output = Dense(units=self._labels_shape['output'], activation=self._last_layer_activation, name='output')(activation_4)
        variance_output = Dense(units=self._labels_shape['output'], activation='linear', name='variance_output')(activation_4)

        model = Model(inputs=[input_tensor, labels_err_tensor], outputs=[output, variance_output])
        # new astroNN high performance dropout variational inference on GPU expects single output
        model_prediction = Model(inputs=[input_tensor], outputs=concatenate([output, variance_output]))

        if self.task == 'regression':
            variance_loss = mse_var_wrapper(output, labels_err_tensor)
            output_loss = mse_lin_wrapper(variance_output, labels_err_tensor)
        elif self.task == 'classification':
            output_loss = bayesian_categorical_crossentropy_wrapper(variance_output)
            variance_loss = bayesian_categorical_crossentropy_var_wrapper(output)
        elif self.task == 'binary_classification':
            output_loss = bayesian_binary_crossentropy_wrapper(variance_output)
            variance_loss = bayesian_binary_crossentropy_var_wrapper(output)
        else:
            raise RuntimeError('Only "regression", "classification" and "binary_classification" are supported')

        return model, model_prediction, output_loss, variance_loss


# noinspection PyCallingNonCallable
class ApogeeBCNNCensored(BayesianCNNBase):
    """
    Class for Bayesian censored convolutional neural network for stellar spectra analysis [specifically APOGEE
    DR14 spectra only]

    Described in the paper: https://ui.adsabs.harvard.edu/abs/2019MNRAS.483.3255L/abstract

    :History: 2018-May-27 - Written - Henry Leung (University of Toronto)
    """

    def __init__(self, lr=0.0005, dropout_rate=0.3):
        super().__init__()

        self._implementation_version = '1.0'
        self.initializer = RandomNormal(mean=0.0, stddev=0.05)
        self.activation = 'relu'
        self.num_filters = [2, 4]
        self.filter_len = 8
        self.pool_length = 4
        # number of neurone for [ApogeeBCNN_Dense_1, ApogeeBCNN_Dense_2, aspcap_1, aspcap_2, hidden]
        self.num_hidden = [192, 96, 32, 16, 2]
        self.max_epochs = 50
        self.lr = lr
        self.reduce_lr_epsilon = 0.00005
        self.maxnorm = .5

        self.reduce_lr_min = 1e-8
        self.reduce_lr_patience = 2
        self.l2 = 5e-9
        self.dropout_rate = dropout_rate

        self.input_norm_mode = 3

        self.task = 'regression'

        self.targetname = ['teff', 'logg', 'C', 'C1', 'N', 'O', 'Na', 'Mg', 'Al', 'Si', 'P', 'S', 'K', 'Ca', 'Ti',
                           'Ti2', 'V', 'Cr', 'Mn', 'Fe', 'Co', 'Ni']

    def model(self):
        input_tensor = Input(shape=self._input_shape['input'], name='input')
        input_tensor_flattened = Flatten()(input_tensor)
        labels_err_tensor = Input(shape=(self._labels_shape['output'],), name='labels_err')

        # slice spectra to censor out useless region for elements

        censored_c_input = BoolMask(aspcap_mask("C", dr=14), name='C_Mask')(input_tensor_flattened)