# Copyright (c) 2021 Graphcore Ltd. All rights reserved.
#
# Keras lstm has a configurable activations, and can be
# exported to onnx, which makes it useful for checking the
# onnx lstm activation attribute. However, I did not want to
# add tensorflow and keras as dependencies of popart just
# for this test. A specific older version of tensorflow
# (2.3.0) is curently required for keras2onnx to work as
# well, making it even less desirable a dependency to add to
# the project.
# This script generates a test file with a number of keras
# models already exported to onnx, with the model results,
# and adds a test for each model.
import numpy as np
from pathlib import Path
import tensorflow as tf

print("TensorFlow version is " + tf.__version__)
import keras2onnx

np.random.seed(0)
tf.random.set_seed(0)


def _generate_keras_lstm(data, lstm_kwargs):
    # Create model
    model = tf.keras.models.Sequential()
    model.add(tf.keras.layers.LSTM(4, **lstm_kwargs))

    # Build model and print summary
    model.build(input_shape=(1, 4, 4))
    model.summary()

    result = model(data)
    result = result.numpy()

    onnx_model = keras2onnx.convert_keras(model)

    return onnx_model, result


def _generate_comparison_test(name, data, lstm_kwargs, expected_activations):
    model, result = _generate_keras_lstm(data, lstm_kwargs)
    return f"""
@pytest.mark.parametrize("lstm_op_pattern", [False, True])
def {name}(lstm_op_pattern):
    data = {repr(data)}
    result = {repr(result)}
    model = {model.SerializeToString()}
    _run_comparison_test(data, result, model, {expected_activations}, lstm_op_pattern)
"""


def _popart_root():
    here = Path(__file__)
    for x in reversed(here.parents):
        if x.name == "popart":
            return x


def main():
    data = np.random.rand(1, 4, 4).astype(np.float32)
    # Add copyright and imports to start of file
    out = f"""# Copyright (c) 2021 Graphcore Ltd. All rights reserved.
#
# THIS IS AN AUTOGENERATED FILE, DO NOT EDIT DIRECTLY
#
# To regenerate this file run:
#     python {Path(__file__).relative_to(_popart_root().parent)}
#
# File generated using TensorFlow version {tf.__version__}
# and keras2onnx version {keras2onnx.__version__}

import numpy as np
from numpy import array, float32
import popart
import onnx
import pytest

# `import test_util` requires adding to sys.path
import sys
from pathlib import Path
sys.path.append(str(Path(__file__).resolve().parent.parent))
import test_util as tu


def _run_comparison_test(data, result, proto, expected_activations, lstm_op_pattern):
    model = onnx.load_from_string(proto)

    if expected_activations:
        lstms = [i for i in model.graph.node if i.op_type == 'LSTM']
        assert len(lstms) == 1
        activations = [i for i in lstms[0].attribute if i.name == 'activations']
        assert len(activations) == 1
        activations = activations[0].strings
        assert len(activations) == len(expected_activations)
        for expected, actual in zip(expected_activations, activations):
            assert expected == actual.decode('utf-8').lower()

    outId = model.graph.output[0].name
    inId = model.graph.input[0].name

    dataFlow = popart.DataFlow(1, {{outId: popart.AnchorReturnType("All")}})
    patterns = popart.Patterns(popart.PatternsLevel.Default)
    patterns.enablePattern('LSTMOp', lstm_op_pattern)
    with tu.create_test_device() as device:
        session = popart.InferenceSession(fnModel=proto,
                                          dataFlow=dataFlow,
                                          deviceInfo=device,
                                          patterns=patterns)

        session.prepareDevice()

        anchors = session.initAnchorArrays()
        stepio = popart.PyStepIO({{inId: data}}, anchors)
        session.run(stepio)

        assert np.allclose(anchors[outId], result)


"""
    # Generate a default keras lstm
    out += _generate_comparison_test("test_basic", data, {}, [])

    keras_supported_activations = ["relu", "sigmoid", "softmax", "tanh"]

    # Generate a test for each combination of activations
    for first in keras_supported_activations:
        for second in keras_supported_activations:
            try:
                out += _generate_comparison_test(
                    f"test_{first}_{second}",
                    data,
                    {"activation": first, "recurrent_activation": second},
                    [second, first, first],
                )
            except NotImplementedError:
                pass

    here = Path(__file__).parent
    out_path = here / "keras_lstm_tests.py"
    out_path.write_text(out, encoding="utf-8")


if __name__ == "__main__":
    main()