# coding: utf-8
import unittest
import numpy as np
from onnx import helper, onnx_pb as onnx_proto
import onnxruntime as _ort
from onnxruntime_extensions import (
    OrtPyFunction,
    PyCustomOpDef,
    onnx_op,
    make_onnx_model,
    get_library_path as _get_library_path,
)


def _create_test_model_segment_sum(prefix, domain="ai.onnx.contrib"):
    nodes = [
        helper.make_node("Identity", ["data"], ["id1"]),
        helper.make_node("Identity", ["segment_ids"], ["id2"]),
        helper.make_node("%sSegmentSum" % prefix, ["id1", "id2"], ["z"], domain=domain),
    ]

    input0 = helper.make_tensor_value_info("data", onnx_proto.TensorProto.FLOAT, [])
    input1 = helper.make_tensor_value_info(
        "segment_ids", onnx_proto.TensorProto.INT64, []
    )
    output0 = helper.make_tensor_value_info("z", onnx_proto.TensorProto.FLOAT, [])

    graph = helper.make_graph(nodes, "test0", [input0, input1], [output0])
    model = make_onnx_model(graph)
    return model


class TestMathOpString(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        @onnx_op(
            op_type="PySegmentSum",
            inputs=[PyCustomOpDef.dt_float, PyCustomOpDef.dt_int64],
            outputs=[PyCustomOpDef.dt_float],
        )
        def segment_sum(data, segment_ids):
            # segment_ids is sorted
            nb_seg = segment_ids[-1] + 1
            sh = (nb_seg,) + data.shape[1:]
            res = np.zeros(sh, dtype=data.dtype)
            for seg, row in zip(segment_ids, data):
                res[seg] += row
            return res

    def test_segment_sum_cc(self):
        so = _ort.SessionOptions()
        so.register_custom_ops_library(_get_library_path())
        onnx_model = _create_test_model_segment_sum("")
        self.assertIn('op_type: "SegmentSum"', str(onnx_model))
        sess = _ort.InferenceSession(onnx_model.SerializeToString(), so)
        data = np.array([[1, 2, 3, 4], [4, 3, 2, 1], [5, 6, 7, 8]], dtype=np.float32)
        segment_ids = np.array([0, 0, 1], dtype=np.int64)
        exp = np.array([[5, 5, 5, 5], [5, 6, 7, 8]], dtype=np.float32)
        txout = sess.run(None, {"data": data, "segment_ids": segment_ids})
        self.assertEqual(exp.shape, txout[0].shape)
        self.assertEqual(exp.tolist(), txout[0].tolist())

    def test_segment_sum_python(self):
        so = _ort.SessionOptions()
        so.register_custom_ops_library(_get_library_path())
        onnx_model = _create_test_model_segment_sum("Py")
        self.assertIn('op_type: "PySegmentSum"', str(onnx_model))
        sess = _ort.InferenceSession(onnx_model.SerializeToString(), so)
        data = np.array([[1, 2, 3, 4], [4, 3, 2, 1], [5, 6, 7, 8]], dtype=np.float32)
        segment_ids = np.array([0, 0, 1], dtype=np.int64)
        exp = np.array([[5, 5, 5, 5], [5, 6, 7, 8]], dtype=np.float32)
        txout = sess.run(None, {"data": data, "segment_ids": segment_ids})
        self.assertEqual(exp.shape, txout[0].shape)
        self.assertEqual(exp.tolist(), txout[0].tolist())

        try:
            from tensorflow.raw_ops import SegmentSum

            dotf = True
        except ImportError:
            dotf = False
        if dotf:
            tfres = SegmentSum(data=data, segment_ids=segment_ids)
            self.assertEqual(tfres.shape, txout[0].shape)
            self.assertEqual(tfres.numpy().tolist(), txout[0].tolist())

    def test_inverse(self):
        mat = np.random.rand(5, 5).astype(np.float32)
        inv_mat = np.linalg.inv(mat)
        ort_inv = OrtPyFunction.from_customop('Inverse')
        act_mat = ort_inv(mat)
        self.assertTrue(np.allclose(inv_mat, act_mat, rtol=1.e-3))


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