from pathlib import Path

import numpy as np
import pytest
import tensorflow as tf
from qkeras.utils import _add_supported_quantized_objects
from sklearn.metrics import accuracy_score
from tensorflow.keras.models import model_from_json

import hls4ml

co = {}
_add_supported_quantized_objects(co)

test_root_path = Path(__file__).parent
example_model_path = (test_root_path / '../../example-models').resolve()


@pytest.fixture(scope='module')
def mnist_data():
    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
    # Scale images to the [0, 1] range
    x_train = x_train.astype('float32') / 255
    x_test = x_test.astype('float32') / 255
    # Make sure images have shape (28, 28, 1)
    x_train = np.expand_dims(x_train, -1)
    x_test = np.expand_dims(x_test, -1)
    y_train = tf.keras.utils.to_categorical(y_train, 10)
    y_test = tf.keras.utils.to_categorical(y_test, 10)
    return x_train, y_train, x_test, y_test


@pytest.fixture(scope='module')
def mnist_model():
    model_path = example_model_path / 'keras/qkeras_mnist_cnn.json'
    with model_path.open('r') as f:
        jsons = f.read()
    model = model_from_json(jsons, custom_objects=co)
    model.load_weights(example_model_path / 'keras/qkeras_mnist_cnn_weights.h5')
    return model


@pytest.fixture
def hls_model(mnist_model, request, test_case_id):
    backend, io_type, strategy = request.param
    keras_model = mnist_model
    hls_config = hls4ml.utils.config_from_keras_model(keras_model, granularity='name', backend=backend)
    hls_config['Model']['Strategy'] = strategy
    hls_config['LayerName']['softmax']['Strategy'] = 'Stable'
    output_dir = str(test_root_path / test_case_id)

    hls_model = hls4ml.converters.convert_from_keras_model(
        keras_model, hls_config=hls_config, output_dir=output_dir, backend=backend, io_type=io_type
    )

    hls_model.compile()
    return hls_model


@pytest.mark.parametrize(
    'hls_model',
    [
        ('Quartus', 'io_parallel', 'resource'),
        ('Quartus', 'io_stream', 'resource'),
        ('Vivado', 'io_parallel', 'resource'),
        ('Vivado', 'io_parallel', 'latency'),
        ('Vivado', 'io_stream', 'latency'),
        ('Vivado', 'io_stream', 'resource'),
        ('Vitis', 'io_parallel', 'resource'),
        ('Vitis', 'io_parallel', 'latency'),
        ('Vitis', 'io_stream', 'latency'),
        ('Vitis', 'io_stream', 'resource'),
    ],
    indirect=True,
    ids=[
        'Quartus_io_parallel_resource',
        'Quartus_io_stream_resource',
        'Vivado_io_parallel_resource',
        'Vivado_io_parallel_latency',
        'Vivado_io_stream_latency',
        'Vivado_io_stream_resource',
        'Vitis_io_parallel_resource',
        'Vitis_io_parallel_latency',
        'Vitis_io_stream_latency',
        'Vitis_io_stream_resource',
    ],
)
def test_accuracy(mnist_data, mnist_model, hls_model):
    x_train, y_train, x_test, y_test = mnist_data
    x_test, y_test = x_test[:1000], y_test[:1000]
    # model under test predictions and accuracy
    y_keras = mnist_model.predict(x_test)
    y_hls4ml = hls_model.predict(x_test)

    acc_keras = accuracy_score(np.argmax(y_test, axis=1), np.argmax(y_keras, axis=1))
    acc_hls4ml = accuracy_score(np.argmax(y_test, axis=1), np.argmax(y_hls4ml, axis=1))
    rel_diff = abs(acc_keras - acc_hls4ml) / acc_keras

    print(f'Accuracy keras:      {acc_keras}')
    print(f'Accuracy hls4ml:     {acc_hls4ml}')
    print(f'Relative difference: {rel_diff}')

    assert acc_keras > 0.92 and rel_diff < 0.01
