#
# Copyright (c) 2021, NVIDIA CORPORATION.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import pytest
import tensorflow as tf
from tensorflow.keras import regularizers

import merlin.models.tf as ml
from merlin.models.data.synthetic import SyntheticData


@pytest.mark.parametrize("dim", [32, 64])
@pytest.mark.parametrize("activation", ["relu", "tanh"])
@pytest.mark.parametrize("dropout", [None, 0.5])
@pytest.mark.parametrize(
    "normalization", [None, "batch_norm", tf.keras.layers.BatchNormalization()]
)
def test_mlp_block_yoochoose(
    testing_data: SyntheticData,
    dim,
    activation,
    dropout,
    normalization,
    kernel_initializer="glorot_uniform",
    bias_initializer="zeros",
    kernel_regularizer="l2",
    bias_regularizer="l1",
    activity_regularizer=regularizers.l2(1e-4),
):
    inputs = ml.InputBlock(testing_data.schema)

    mlp = ml.MLPBlock(
        [dim],
        activation=activation,
        dropout=dropout,
        normalization=normalization,
        kernel_initializer=kernel_initializer,
        bias_initializer=bias_initializer,
        kernel_regularizer=kernel_regularizer,
        bias_regularizer=bias_regularizer,
        activity_regularizer=activity_regularizer,
    )
    body = ml.SequentialBlock([inputs, mlp])

    outputs = body(testing_data.tf_tensor_dict)

    assert list(outputs.shape) == [100, dim]
    assert mlp.layers[0].units == dim
    assert mlp.layers[0].dense.activation.__name__ == activation
    if dropout:
        assert mlp.layers[1].rate == dropout
    if normalization:
        if normalization == "batch_norm":
            normalization = tf.keras.layers.BatchNormalization()

        assert mlp.layers[-1].__class__.__name__ == normalization.__class__.__name__
