import os

import numpy as np
import onnx
import oyaml as yaml
import pytest

import audeer
import audobject

import audonnx.testing


def feature_addition(feature, offset=1):
    return feature + offset


def min_max(x, sr):
    import numpy as np

    return np.array([x.min(), x.max()], np.float32)


@pytest.mark.parametrize(
    "object, transform, labels, expected",
    [
        (
            audonnx.testing.create_model_proto([[2]]),
            audonnx.Function(min_max),
            {"output-0": ["min", "max"]},
            min_max(pytest.SIGNAL, pytest.SAMPLING_RATE),
        ),
    ],
)
def test_load_legacy(tmpdir, object, transform, labels, expected):
    root = tmpdir

    onnx_path = os.path.join(root, "model.onnx")
    onnx.save(object, onnx_path)

    # create from onnx

    transform_path = os.path.join(root, "transform.yaml")
    transform.to_yaml(transform_path)

    labels_path = os.path.join(root, "labels.yaml")
    with open(labels_path, "w") as fp:
        yaml.dump(labels, fp)

    model = audonnx.load(
        root,
        model_file=os.path.basename(onnx_path),
        transform_file=os.path.basename(transform_path),
        labels_file=os.path.basename(labels_path),
    )
    y = model(pytest.SIGNAL, pytest.SAMPLING_RATE)

    np.testing.assert_almost_equal(y, expected, decimal=1)
    for key, values in labels.items():
        assert model.outputs[key].labels == labels[key]

    # legacy mode -> load from ONNX if YAML does not exist

    model = audonnx.load(
        root,
        model_file=audeer.replace_file_extension(onnx_path, "yaml"),
    )
    y = model(pytest.SIGNAL, pytest.SAMPLING_RATE)

    np.testing.assert_almost_equal(y, expected, decimal=1)
    for key, values in labels.items():
        assert model.outputs[key].labels == labels[key]

    # store wrong model YAML

    model_path = os.path.join(tmpdir, "single.yaml")
    audobject.Object().to_yaml(model_path)
    model = audonnx.load(
        root,
        model_file=os.path.basename(onnx_path),
    )

    # file extension does not end on '.yaml'

    model_path = os.path.join(tmpdir, "model.yml")
    with pytest.raises(ValueError):
        model.to_yaml(model_path)

    # create from YAML

    model_path = os.path.join(tmpdir, "model.yaml")
    model.to_yaml(model_path)

    model = audonnx.load(tmpdir)
    y = model(pytest.SIGNAL, pytest.SAMPLING_RATE)

    np.testing.assert_equal(y, expected)
    for key, values in labels.items():
        assert model.outputs[key].labels == labels[key]


@pytest.mark.parametrize(
    "model, inputs, sampling_rate, expected",
    [
        # Feature transformed with additional argument
        (
            audonnx.Model(
                audonnx.testing.create_model_proto([[2]]),
                transform={"input-0": audonnx.Function(feature_addition)},
            ),
            {"feature": np.array([1.0, 2.0], dtype=np.float32), "offset": 2},
            pytest.SAMPLING_RATE,
            np.array([3.0, 4.0], dtype=np.float32),
        ),
    ],
)
def test_load_custom_signature(tmpdir, model, inputs, sampling_rate, expected):
    yaml_path = audeer.path(tmpdir, "model.yaml")
    model.to_yaml(yaml_path)

    # Load from onnx and check that results are as expected
    loaded_model = audonnx.load(tmpdir)

    y = loaded_model(inputs, sampling_rate)
    if isinstance(y, dict):
        for key, values in y.items():
            np.testing.assert_equal(values, expected[key])
    else:
        np.testing.assert_equal(y, expected)
