import unittest
import numpy as np
from onnx import helper, onnx_pb as onnx_proto
import onnxruntime as _ort
from onnxruntime_extensions import get_library_path as _get_library_path, make_onnx_model


def _create_test_model(input_dims, output_dims):
    nodes = []
    nodes[0:] = [helper.make_node('Identity', ['input_1'], ['left'])]
    nodes[1:] = [helper.make_node('Identity', ['input_2'], ['right'])]
    nodes[2:] = [helper.make_node(
        'StringConcat', ['left', 'right'], ['output'], domain='ai.onnx.contrib')]

    input1 = helper.make_tensor_value_info(
        'input_1', onnx_proto.TensorProto.STRING, [None] * input_dims)
    input2 = helper.make_tensor_value_info(
        'input_2', onnx_proto.TensorProto.STRING, [None] * input_dims)
    output = helper.make_tensor_value_info(
        'output', onnx_proto.TensorProto.STRING, [None] * output_dims)

    graph = helper.make_graph(nodes, 'test0', [input1, input2], [output])
    model = make_onnx_model(graph, 12)
    return model


def _run_string_concat(input1, input2):
    model = _create_test_model(input1.ndim, input1.ndim)

    so = _ort.SessionOptions()
    so.register_custom_ops_library(_get_library_path())
    sess = _ort.InferenceSession(model.SerializeToString(), so, providers=['CPUExecutionProvider'])
    result = sess.run(None, {'input_1': input1, 'input_2': input2})

    # verify
    output = []
    shape = input1.shape
    input1 = input1.flatten()
    input2 = input2.flatten()
    for i in range(len(input1)):
        output.append(input1[i] + input2[i])
    output = np.array(output).reshape(shape)
    np.testing.assert_array_equal(result, [output])


class TestStringConcat(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        pass

    def test_string_concat(self):
        _run_string_concat(np.array(["a"]), np.array(["b"]))
        _run_string_concat(np.array(["a", "b", "c", "d"]), np.array(["d", "c", "b", "a"]))
        _run_string_concat(np.array([["a", "b"], ["c", "d"]]), np.array([["d", "c"], ["b", "a"]]))
        _run_string_concat(np.array(["你好"]), np.array(["这是一个测试"]))
        _run_string_concat(np.array(["すみません"]), np.array(["これはテストです"]))
        _run_string_concat(np.array(["👾 🤖 🎃 😺 😸 😹"]), np.array(["😻 😼 😽 🙀 😿 😾"]))
        _run_string_concat(np.array(["龖"]), np.array(["龘讋"]))


if __name__ == "__main__":
    unittest.main()