import keras

from typing import Tuple

from keras import backend as K

from keras.models import Model

from keras.layers import Input, BatchNormalization, Activation

from keras.layers.core import Dense, Dropout

from keras.optimizers import Adam

from keras.regularizers import l1

def baseline_autoencoder(

    input_shape: Tuple, dropout: float, num_ae_units1: int, num_ae_units2: int, l1_reg_lambda_exp: float,

) -> Model:

    """Baseline autoencoder as published in https://www.nature.com/articles/s42256-019-0019-2"""

    inputvec = Input(shape=(input_shape,))

    x = Dropout(dropout, input_shape=(input_shape,))(inputvec)

    x = Dense(units=int(num_ae_units1), activation='relu', activity_regularizer=l1(10**l1_reg_lambda_exp))(x)

    encoded = Dense(units=int(num_ae_units2), activation='relu', name='encoded')(x)

    risk_pred = Dense(units=1,  activation='linear', name='predicted_risk')(encoded)

    z = Dense(units=int(num_ae_units1), activation='relu')(encoded)

    decoded = Dense(units=input_shape, activation='linear', name='decoded')(z)

    model = Model(inputs=inputvec, outputs=[decoded, risk_pred])

    return model

def baseline_bn_autoencoder(

    input_shape: Tuple, dropout: float, num_ae_units1: int, num_ae_units2: int, l1_reg_lambda_exp: float,

) -> Model:

    """Add batch normalization to each layer before relu activation, based on baseline_autoencoder."""

    inputvec = Input(shape=(input_shape,))

    x = Dropout(dropout, input_shape=(input_shape,))(inputvec)

    x = Dense(units=int(num_ae_units1), activation=None, activity_regularizer=l1(10**l1_reg_lambda_exp))(x)

    x = BatchNormalization()(x)

    x = Activation("relu")(x)

    x = Dense(units=int(num_ae_units2), activation=None)(x)

    x = BatchNormalization()(x)

    encoded = Activation("relu", name='encoded')(x)

    risk_pred = Dense(units=1,  activation='linear', name='predicted_risk')(encoded)

    x = Dense(units=int(num_ae_units1), activation=None)(encoded)

    x = BatchNormalization()(x)

    z = Activation("relu")(x)

    decoded = Dense(units=input_shape, activation='linear', name='decoded')(z)

    model = Model(inputs=inputvec, outputs=[decoded, risk_pred])

    return model

def model3_bn_autoencoder(

    input_shape: Tuple, dropout: float, num_ae_units1: int, num_ae_units2: int, num_risk_units: int,

    l1_reg_lambda_exp: float,

) -> Model:

    """

    Add one more relu layer between encoded and risk_pred, based on baseline_bn_autoencoder.

    Model 3 architecture: https://arxiv.org/pdf/1910.02951v1.pdf

    """

    inputvec = Input(shape=(input_shape,))

    x = Dropout(dropout, input_shape=(input_shape,))(inputvec)

    x = Dense(units=int(num_ae_units1), activation=None, activity_regularizer=l1(10**l1_reg_lambda_exp))(x)

    x = BatchNormalization()(x)

    x = Activation("relu")(x)

    x = Dense(units=int(num_ae_units2), activation=None)(x)

    x = BatchNormalization()(x)

    encoded = Activation("relu", name='encoded')(x)

    x = Dense(units=num_risk_units,  activation=None)(encoded)

    x = BatchNormalization()(x)

    x = Activation("relu", name='encoded')(x)

    risk_pred = Dense(units=1,  activation='linear', name='predicted_risk')(x)

    x = Dense(units=int(num_ae_units1), activation=None)(encoded)

    x = BatchNormalization()(x)

    z = Activation("relu")(x)

    decoded = Dense(units=input_shape, activation='linear', name='decoded')(z)

    model = Model(inputs=inputvec, outputs=[decoded, risk_pred])

    return model

def deep_model3_bn_autoencoder(

    input_shape: Tuple, dropout: float, num_ae_units1: int, num_ae_units2: int, num_ae_units3: int,

    num_risk_units: int, l1_reg_lambda_exp: float,

) -> Model:

    """

    Add one more relu layer in autoencoder, based on model3_bn_autoencoder.

    Model 3 architecture: https://arxiv.org/pdf/1910.02951v1.pdf

    """

    inputvec = Input(shape=(input_shape,))

    x = Dropout(dropout, input_shape=(input_shape,))(inputvec)

    x = Dense(units=int(num_ae_units1), activation=None, activity_regularizer=l1(10**l1_reg_lambda_exp))(x)

    x = BatchNormalization()(x)

    x = Activation("relu")(x)

    x = Dense(units=int(num_ae_units2), activation=None)(x)

    x = BatchNormalization()(x)

    x = Activation("relu")(x)

    x = Dense(units=int(num_ae_units3), activation=None)(x)

    x = BatchNormalization()(x)

    encoded = Activation("relu", name='encoded')(x)

    x = Dense(units=num_risk_units,  activation=None)(encoded)

    x = BatchNormalization()(x)

    x = Activation("relu", name='encoded')(x)

    risk_pred = Dense(units=1,  activation='linear', name='predicted_risk')(x)

    x = Dense(units=int(num_ae_units2), activation=None)(encoded)

    x = BatchNormalization()(x)

    x = Activation("relu")(x)

    x = Dense(units=int(num_ae_units1), activation=None)(x)

    x = BatchNormalization()(x)

    z = Activation("relu")(x)

    decoded = Dense(units=input_shape, activation='linear', name='decoded')(z)

    model = Model(inputs=inputvec, outputs=[decoded, risk_pred])

    return model

def model_factory(model_name: str, **kwargs):

    # Before defining network architecture, clear current computation graph (if one exists)

    K.clear_session()

    if model_name == "baseline_autoencoder":

        model = baseline_autoencoder(**kwargs)

    elif model_name == "baseline_bn_autoencoder":

        model = baseline_bn_autoencoder(**kwargs)

    elif model_name == "model3_bn_autoencoder":

        model = model3_bn_autoencoder(**kwargs)

    elif model_name == "deep_model3_bn_autoencoder":

        model = deep_model3_bn_autoencoder(**kwargs)

    else:

        raise ValueError("Model name {} has not been implemented.".format(model_name))

    return model