from pathlib import Path

import numpy as np
import pytest
import tensorflow as tf

import hls4ml

test_root_path = Path(__file__).parent

padds_options = ['same', 'valid']
chans_options = ['channels_last']
strides_options = [(1, 1), (2, 2)]
kernel_options = [(2, 2), (3, 3)]
bias_options = [False]
rf_options = [1, 3, 4]  # each rf corresponds to one of the three cases of depthwise resource for io_stream
input_size_options = [3]


@pytest.mark.parametrize('chans', chans_options)
@pytest.mark.parametrize('padds', padds_options)
@pytest.mark.parametrize('strides', strides_options)
@pytest.mark.parametrize('kernels', kernel_options)
@pytest.mark.parametrize('bias', bias_options)
@pytest.mark.parametrize(
    'backend, io_type, strategy',
    [
        ('Vivado', 'io_parallel', 'latency'),
        ('Vitis', 'io_parallel', 'latency'),
        ('Vivado', 'io_stream', 'latency'),
        ('Vitis', 'io_stream', 'latency'),
        ('Vivado', 'io_stream', 'resource'),
        ('Vitis', 'io_stream', 'resource'),
        ('Catapult', 'io_stream', 'latency'),
    ],
)
@pytest.mark.parametrize('rf', rf_options)
@pytest.mark.parametrize('input_size', input_size_options)
def test_sepconv2d(test_case_id, chans, padds, strides, kernels, bias, io_type, backend, strategy, rf, input_size):
    model = tf.keras.models.Sequential()
    input_shape = (16, 16, input_size)
    model.add(
        tf.keras.layers.SeparableConv2D(
            filters=8,
            kernel_size=kernels,
            strides=strides,
            padding=padds,
            input_shape=input_shape,
            kernel_initializer='normal',
            use_bias=bias,
            data_format=chans,
        )
    )

    model.compile(optimizer='adam', loss='mse')
    X_input = np.random.rand(100, *input_shape)
    keras_prediction = model.predict(X_input)
    config = hls4ml.utils.config_from_keras_model(
        model, default_precision='ap_fixed<32,8>', granularity='name', backend=backend
    )
    config['Model']['Strategy'] = strategy
    config['Model']['ReuseFactor'] = rf
    output_dir = str(test_root_path / test_case_id)
    hls_model = hls4ml.converters.convert_from_keras_model(
        model, hls_config=config, output_dir=output_dir, io_type=io_type, backend=backend
    )
    hls_model.compile()
    hls_prediction = hls_model.predict(X_input).reshape(keras_prediction.shape)

    np.testing.assert_allclose(hls_prediction, keras_prediction, rtol=0, atol=0.001)
