"""Regression with abstention network architecture."""

import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import Dense
from tensorflow.keras import regularizers
from tensorflow.keras.models import Sequential

__author__ = "Elizabeth A. Barnes and Randal J. Barnes"
__date__ = "March 4, 2021"

np.warnings.filterwarnings('ignore', category=np.VisibleDeprecationWarning)


def defineNN(hiddens, input_shape, output_shape, ridge_penalty=0., act_fun='relu', network_seed=99):

    inputs = tf.keras.Input( shape = input_shape )
    x = inputs
    
    # linear network only
    if hiddens[0] == 0:
        
        x = tf.keras.layers.Dense(1,
                                  activation='linear',
                                  use_bias=True,
                                  kernel_regularizer=regularizers.l1_l2(l1=0.00, l2=ridge_penalty),
                                  bias_initializer=tf.keras.initializers.RandomNormal(seed=network_seed),
                                  kernel_initializer=tf.keras.initializers.RandomNormal(seed=network_seed),
                                 )(x)
    else:
        # initialize first layer
        x = tf.keras.layers.Dense(hiddens[0],
                                  activation=act_fun,
                                  use_bias=True,
                                  kernel_regularizer=regularizers.l1_l2(l1=0.00, l2=ridge_penalty),
                                  bias_initializer=tf.keras.initializers.RandomNormal(seed=network_seed),
                                  kernel_initializer=tf.keras.initializers.RandomNormal(seed=network_seed),                          
                                 )(x)
        
        # initialize other layers
        for layer in hiddens[1:]:
            x = tf.keras.layers.Dense(layer,
                                      activation=act_fun,
                                      use_bias=True,
                                      kernel_regularizer=regularizers.l1_l2(l1=0.00, l2=0.0),
                                      bias_initializer=tf.keras.initializers.RandomNormal(seed=network_seed),
                                      kernel_initializer=tf.keras.initializers.RandomNormal(seed=network_seed),                          
                                     )(x)  
            
    # set final output units separately
    mu_unit = tf.keras.layers.Dense(1, 
                                    activation = 'linear',
                                    use_bias=True,
                                    bias_initializer=tf.keras.initializers.RandomNormal(seed=network_seed),
                                    kernel_initializer=tf.keras.initializers.RandomNormal(seed=network_seed),                      
#                                     bias_initializer=tf.keras.initializers.RandomUniform(minval=0.,maxval=1.,seed=network_seed),   
#                                     kernel_initializer=tf.keras.initializers.RandomUniform(minval=0.,maxval=1.,seed=network_seed),
                                   )(x)
    
    sigma_unit = tf.keras.layers.Dense(1, 
                                    activation = 'relu',       # constrain sigma to be positive
                                    use_bias=True,
                                    bias_initializer=tf.keras.initializers.RandomUniform(minval=.1,maxval= .5,seed=network_seed),   
                                    kernel_initializer=tf.keras.initializers.RandomUniform(minval=.1,maxval=.5,seed=network_seed),
#                                     kernel_constraint = tf.keras.constraints.NonNeg(),	                                       
#                                     bias_constraint = tf.keras.constraints.NonNeg(),	                                       
                                   )(x)    

    # final output layer
    output_layer = tf.keras.layers.concatenate([mu_unit,sigma_unit],axis=1)
    
    # finalize the model
    model = tf.keras.models.Model(inputs=inputs, outputs=output_layer)            
    
    return model


tensorflow.keras.Input tensorflow.keras.initializers.RandomUniform tensorflow.keras.models.Model tensorflow.keras.regularizers.l1_l2 tensorflow.keras.layers.concatenate numpy.warnings.filterwarnings tensorflow.keras.initializers.RandomNormal
[(12, 'numpy.warnings.filterwarnings', 'np.warnings.filterwarnings', (['"""ignore"""'], {'category': 'np.VisibleDeprecationWarning'}), True, 'import numpy as np\n'), (17, 'tensorflow.keras.Input', 'tf.keras.Input', ([], {'shape': 'input_shape'}), True, 'import tensorflow as tf\n'), (70, 'tensorflow.keras.layers.concatenate', 'tf.keras.layers.concatenate', (['[mu_unit, sigma_unit]'], {'axis': '(1)'}), True, 'import tensorflow as tf\n'), (73, 'tensorflow.keras.models.Model', 'tf.keras.models.Model', ([], {'inputs': 'inputs', 'outputs': 'output_layer'}), True, 'import tensorflow as tf\n'), (54, 'tensorflow.keras.initializers.RandomNormal', 'tf.keras.initializers.RandomNormal', ([], {'seed': 'network_seed'}), True, 'import tensorflow as tf\n'), (55, 'tensorflow.keras.initializers.RandomNormal', 'tf.keras.initializers.RandomNormal', ([], {'seed': 'network_seed'}), True, 'import tensorflow as tf\n'), (63, 'tensorflow.keras.initializers.RandomUniform', 'tf.keras.initializers.RandomUniform', ([], {'minval': '(0.1)', 'maxval': '(0.5)', 'seed': 'network_seed'}), True, 'import tensorflow as tf\n'), (64, 'tensorflow.keras.initializers.RandomUniform', 'tf.keras.initializers.RandomUniform', ([], {'minval': '(0.1)', 'maxval': '(0.5)', 'seed': 'network_seed'}), True, 'import tensorflow as tf\n'), (26, 'tensorflow.keras.regularizers.l1_l2', 'regularizers.l1_l2', ([], {'l1': '(0.0)', 'l2': 'ridge_penalty'}), False, 'from tensorflow.keras import regularizers\n'), (27, 'tensorflow.keras.initializers.RandomNormal', 'tf.keras.initializers.RandomNormal', ([], {'seed': 'network_seed'}), True, 'import tensorflow as tf\n'), (28, 'tensorflow.keras.initializers.RandomNormal', 'tf.keras.initializers.RandomNormal', ([], {'seed': 'network_seed'}), True, 'import tensorflow as tf\n'), (35, 'tensorflow.keras.regularizers.l1_l2', 'regularizers.l1_l2', ([], {'l1': '(0.0)', 'l2': 'ridge_penalty'}), False, 'from tensorflow.keras import regularizers\n'), (36, 'tensorflow.keras.initializers.RandomNormal', 'tf.keras.initializers.RandomNormal', ([], {'seed': 'network_seed'}), True, 'import tensorflow as tf\n'), (37, 'tensorflow.keras.initializers.RandomNormal', 'tf.keras.initializers.RandomNormal', ([], {'seed': 'network_seed'}), True, 'import tensorflow as tf\n'), (45, 'tensorflow.keras.regularizers.l1_l2', 'regularizers.l1_l2', ([], {'l1': '(0.0)', 'l2': '(0.0)'}), False, 'from tensorflow.keras import regularizers\n'), (46, 'tensorflow.keras.initializers.RandomNormal', 'tf.keras.initializers.RandomNormal', ([], {'seed': 'network_seed'}), True, 'import tensorflow as tf\n'), (47, 'tensorflow.keras.initializers.RandomNormal', 'tf.keras.initializers.RandomNormal', ([], {'seed': 'network_seed'}), True, 'import tensorflow as tf\n')]