# Path: mimic3models/keras_utils.py
# class DecompensationMetrics(keras.callbacks.Callback):
# class InHospitalMortalityMetrics(keras.callbacks.Callback):
# class PhenotypingMetrics(keras.callbacks.Callback):
# class LengthOfStayMetrics(keras.callbacks.Callback):
# class MultitaskMetrics(keras.callbacks.Callback):
# class CollectAttetion(Layer):
# class Slice(Layer):
# class GetTimestep(Layer):
# class ExtendMask(Layer):
#     def __init__(self, train_data_gen, val_data_gen, deep_supervision,
#                  batch_size=32, early_stopping=True, verbose=2):
#     def calc_metrics(self, data_gen, history, dataset, logs):
#     def on_epoch_end(self, epoch, logs={}):
#     def __init__(self, train_data, val_data, target_repl, batch_size=32, early_stopping=True, verbose=2):
#     def calc_metrics(self, data, history, dataset, logs):
#     def on_epoch_end(self, epoch, logs={}):
#     def __init__(self, train_data_gen, val_data_gen, batch_size=32,
#                  early_stopping=True, verbose=2):
#     def calc_metrics(self, data_gen, history, dataset, logs):
#     def on_epoch_end(self, epoch, logs={}):
#     def __init__(self, train_data_gen, val_data_gen, partition, batch_size=32,
#                  early_stopping=True, verbose=2):
#     def calc_metrics(self, data_gen, history, dataset, logs):
#     def on_epoch_end(self, epoch, logs={}):
#     def __init__(self, train_data_gen, val_data_gen, partition,
#                  batch_size=32, early_stopping=True, verbose=2):
#     def calc_metrics(self, data_gen, history, dataset, logs):
#     def on_epoch_end(self, epoch, logs={}):
# def softmax(x, axis, mask=None):
# def _collect_attention(x, a, mask):
#     def __init__(self, **kwargs):
#     def call(self, inputs, mask=None):
#     def compute_output_shape(self, input_shape):
#     def compute_mask(self, input, input_mask=None):
#     def __init__(self, indices, **kwargs):
#     def call(self, x, mask=None):
#     def compute_output_shape(self, input_shape):
#     def compute_mask(self, input, input_mask=None):
#     def get_config(self):
#     def __init__(self, pos=-1, **kwargs):
#     def call(self, x, mask=None):
#     def compute_output_shape(self, input_shape):
#     def compute_mask(self, input, input_mask=None):
#     def get_config(self):
#     def __init__(self, add_epsilon=False, **kwargs):
#     def call(self, x, mask=None):
#     def compute_output_shape(self, input_shape):
#     def compute_mask(self, input, input_mask=None):
#     def get_config(self):
#         B = self.batch_size


from keras.models import Model
from keras.layers import Input, Dense, LSTM, Masking, Dropout
from keras.layers.wrappers import TimeDistributed
from mimic3models.keras_utils import Slice, GetTimestep, LastTimestep, ExtendMask
from keras.layers.merge import Concatenate, Multiply
        #   if ihm_M = 1 then we will calculate an error term
        #   if ihm_M = 0, our prediction will be 0 and as the label
        #   will also be 0 then error_term will be 0.
        if target_repl:
            ihm_seq = TimeDistributed(Dense(1, activation='sigmoid'), name='ihm_seq')(L)
            ihm_y = GetTimestep(ihm_pos)(ihm_seq)
            ihm_y = Multiply(name='ihm_single')([ihm_y, ihm_M])
            outputs += [ihm_y, ihm_seq]
        else:
            ihm_seq = TimeDistributed(Dense(1, activation='sigmoid'))(L)
            ihm_y = GetTimestep(ihm_pos)(ihm_seq)
            ihm_y = Multiply(name='ihm')([ihm_y, ihm_M])
            outputs += [ihm_y]

        # decomp output
        decomp_y = TimeDistributed(Dense(1, activation='sigmoid'))(L)
        decomp_y = ExtendMask(name='decomp', add_epsilon=True)([decomp_y, decomp_M])
        outputs += [decomp_y]

        # los output
        if partition == 'none':
            los_y = TimeDistributed(Dense(1, activation='relu'))(L)
        else:
            los_y = TimeDistributed(Dense(10, activation='softmax'))(L)
        los_y = ExtendMask(name='los', add_epsilon=True)([los_y, los_M])
        outputs += [los_y]

        # pheno output
        if target_repl:
            pheno_seq = TimeDistributed(Dense(25, activation='sigmoid'), name='pheno_seq')(L)

# Path: mimic3models/keras_utils.py
# class DecompensationMetrics(keras.callbacks.Callback):
# class InHospitalMortalityMetrics(keras.callbacks.Callback):
# class PhenotypingMetrics(keras.callbacks.Callback):
# class LengthOfStayMetrics(keras.callbacks.Callback):
# class MultitaskMetrics(keras.callbacks.Callback):
# class CollectAttetion(Layer):
# class Slice(Layer):
# class GetTimestep(Layer):
# class ExtendMask(Layer):
#     def __init__(self, train_data_gen, val_data_gen, deep_supervision,
#                  batch_size=32, early_stopping=True, verbose=2):
#     def calc_metrics(self, data_gen, history, dataset, logs):
#     def on_epoch_end(self, epoch, logs={}):
#     def __init__(self, train_data, val_data, target_repl, batch_size=32, early_stopping=True, verbose=2):
#     def calc_metrics(self, data, history, dataset, logs):
#     def on_epoch_end(self, epoch, logs={}):
#     def __init__(self, train_data_gen, val_data_gen, batch_size=32,
#                  early_stopping=True, verbose=2):
#     def calc_metrics(self, data_gen, history, dataset, logs):
#     def on_epoch_end(self, epoch, logs={}):
#     def __init__(self, train_data_gen, val_data_gen, partition, batch_size=32,
#                  early_stopping=True, verbose=2):
#     def calc_metrics(self, data_gen, history, dataset, logs):
#     def on_epoch_end(self, epoch, logs={}):
#     def __init__(self, train_data_gen, val_data_gen, partition,
#                  batch_size=32, early_stopping=True, verbose=2):
#     def calc_metrics(self, data_gen, history, dataset, logs):
#     def on_epoch_end(self, epoch, logs={}):
# def softmax(x, axis, mask=None):
# def _collect_attention(x, a, mask):
#     def __init__(self, **kwargs):
#     def call(self, inputs, mask=None):
#     def compute_output_shape(self, input_shape):
#     def compute_mask(self, input, input_mask=None):
#     def __init__(self, indices, **kwargs):
#     def call(self, x, mask=None):
#     def compute_output_shape(self, input_shape):
#     def compute_mask(self, input, input_mask=None):
#     def get_config(self):
#     def __init__(self, pos=-1, **kwargs):
#     def call(self, x, mask=None):
#     def compute_output_shape(self, input_shape):
#     def compute_mask(self, input, input_mask=None):
#     def get_config(self):
#     def __init__(self, add_epsilon=False, **kwargs):
#     def call(self, x, mask=None):
#     def compute_output_shape(self, input_shape):
#     def compute_mask(self, input, input_mask=None):
#     def get_config(self):
#         B = self.batch_size


# Path: mimic3models/keras_models/multitask_channel_wise_lstms.py
from keras.models import Model
from keras.layers import Input, Dense, LSTM, Masking, Dropout
from keras.layers.wrappers import TimeDistributed
from mimic3models.keras_utils import Slice, GetTimestep, LastTimestep, ExtendMask
from keras.layers.merge import Concatenate, Multiply

        #   if ihm_M = 1 then we will calculate an error term
        #   if ihm_M = 0, our prediction will be 0 and as the label
        #   will also be 0 then error_term will be 0.
        if target_repl:
            ihm_seq = TimeDistributed(Dense(1, activation='sigmoid'), name='ihm_seq')(L)
            ihm_y = GetTimestep(ihm_pos)(ihm_seq)
            ihm_y = Multiply(name='ihm_single')([ihm_y, ihm_M])
            outputs += [ihm_y, ihm_seq]
        else:
            ihm_seq = TimeDistributed(Dense(1, activation='sigmoid'))(L)
            ihm_y = GetTimestep(ihm_pos)(ihm_seq)
            ihm_y = Multiply(name='ihm')([ihm_y, ihm_M])
            outputs += [ihm_y]

        # decomp output
        decomp_y = TimeDistributed(Dense(1, activation='sigmoid'))(L)
        decomp_y = ExtendMask(name='decomp', add_epsilon=True)([decomp_y, decomp_M])
        outputs += [decomp_y]

        # los output
        if partition == 'none':
            los_y = TimeDistributed(Dense(1, activation='relu'))(L)
        else:
            los_y = TimeDistributed(Dense(10, activation='softmax'))(L)
        los_y = ExtendMask(name='los', add_epsilon=True)([los_y, los_M])
        outputs += [los_y]

        # pheno output
        if target_repl:
            pheno_seq = TimeDistributed(Dense(25, activation='sigmoid'), name='pheno_seq')(L)

            pheno_y = LastTimestep(name='pheno_single')(pheno_seq)