# import tensorflow as tf
from tensorflow.keras import backend as K


def _bn_relu(layer, dropout=0, **params):
    from tensorflow.keras.layers import BatchNormalization
    from tensorflow.keras.layers import Activation
    layer = BatchNormalization()(layer)
    layer = Activation(params["conv_activation"])(layer)

    if dropout > 0:
        from tensorflow.keras.layers import Dropout
        layer = Dropout(params["conv_dropout"])(layer)

    return layer


def add_conv_weight(
        layer,
        filter_length,
        num_filters,
        subsample_length=1,
        **params):
    from tensorflow.keras.layers import Conv1D
    layer = Conv1D(
        filters=num_filters,
        kernel_size=filter_length,
        strides=subsample_length,
        padding='same',
        kernel_initializer=params["conv_init"])(layer)
    return layer


def add_conv_layers(layer, **params):
    for subsample_length in params["conv_subsample_lengths"]:
        layer = add_conv_weight(
            layer,
            params["conv_filter_length"],
            params["conv_num_filters_start"],
            subsample_length=subsample_length,
            **params)
        layer = _bn_relu(layer, **params)
    return layer


def resnet_block(
        layer,
        num_filters,
        subsample_length,
        block_index,
        **params):
    from tensorflow.keras.layers import Add
    from tensorflow.keras.layers import MaxPooling1D
    from tensorflow.keras.layers import Lambda

    def zeropad(x):
        y = K.zeros_like(x)
        return K.concatenate([x, y], axis=2)

    def zeropad_output_shape(input_shape):
        shape = list(input_shape)
        assert len(shape) == 3
        shape[2] *= 2
        return tuple(shape)

    shortcut = MaxPooling1D(pool_size=subsample_length)(layer)
    zero_pad = (block_index % params["conv_increase_channels_at"]) == 0 \
               and block_index > 0
    if zero_pad is True:
        shortcut = Lambda(zeropad, output_shape=zeropad_output_shape)(shortcut)

    for i in range(params["conv_num_skip"]):
        if not (block_index == 0 and i == 0):
            layer = _bn_relu(
                layer,
                dropout=params["conv_dropout"] if i > 0 else 0,
                **params)
        layer = add_conv_weight(
            layer,
            params["conv_filter_length"],
            num_filters,
            subsample_length if i == 0 else 1,
            **params)
    layer = Add()([shortcut, layer])
    return layer


def get_num_filters_at_index(index, num_start_filters, **params):
    return 2 ** int(index / params["conv_increase_channels_at"]) \
           * num_start_filters


def add_resnet_layers(layer, **params):
    layer = add_conv_weight(
        layer,
        params["conv_filter_length"],
        params["conv_num_filters_start"],
        subsample_length=1,
        **params)
    layer = _bn_relu(layer, **params)
    for index, subsample_length in enumerate(params["conv_subsample_lengths"]):
        num_filters = get_num_filters_at_index(
            index, params["conv_num_filters_start"], **params)
        layer = resnet_block(
            layer,
            num_filters,
            subsample_length,
            index,
            **params)
    layer = _bn_relu(layer, **params)
    return layer


def add_output_layer(layer, **params):
    from tensorflow.keras.layers import Dense, Activation
    from tensorflow.keras.layers import TimeDistributed
    layer = TimeDistributed(Dense(params["num_categories"], name='output_app_dense'),
                            name='output_app_time_distributed')(layer)
    return Activation('softmax', name='app_output')(layer)


def add_compile(model, **params):
    from tensorflow.keras.optimizers import Adam
    optimizer = Adam(
        lr=params["learning_rate"],
        clipnorm=params.get("clipnorm", 1))

    model.compile(loss='categorical_crossentropy',
                  optimizer=optimizer,
                  metrics=['accuracy'])


def build_network(**params):
    from tensorflow.keras.models import Model
    from tensorflow.keras.layers import Input
    inputs = Input(shape=params['input_shape'],
                   dtype='float32',
                   name='inputs')

    if params.get('is_regular_conv', False):
        layer = add_conv_layers(inputs, **params)
    else:
        layer = add_resnet_layers(inputs, **params)

    output = add_output_layer(layer, **params)
    model = Model(inputs=[inputs], outputs=[output], name='Model_app')
    model(inputs)
    if params['model_weights_path'] != "":
        print("loading weights of:{}".format(params['model_weights_path']))
        model.load_weights(params['model_weights_path'])

    if params.get("compile", False):
        add_compile(model, **params)
    return model
