# SPDX-License-Identifier: Apache-2.0

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import numpy as np  # type: ignore

import onnx
import typing
from ..base import Base
from . import expect
from onnx import TensorProto
from typing import List, Optional, Text, Union


def SequenceEmptyImpl() -> List[Optional[np.ndarray]]:
    return []


def SequenceConstructImpl(*tensors: np.ndarray) -> List[np.ndarray]:
    return list(tensors)


def SequenceInsertImpl(sequence: List[np.ndarray], tensor: np.ndarray, position: Optional[int] = None) -> List[np.ndarray]:
    if position is None:
        position = len(sequence)
    sequence.insert(position, tensor)
    return sequence


def SequenceAtImpl(sequence: List[np.ndarray], position: int) -> np.ndarray:
    return sequence[position]


def SequenceEraseImpl(sequence: List[np.ndarray], position: Optional[int] = None) -> List[Optional[np.ndarray]]:
    if position is None:
        position = -1
    del sequence[position]
    return sequence


def SequenceLengthImpl(sequence: List[np.ndarray]) -> np.int64:
    return np.int64(len(sequence))


def SplitToSequenceImpl(tensor: np.ndarray, split: Optional[Union[int, List[int]]] = None, axis: int = 0, keepdims: int = 1) -> List[np.ndarray]:
    dim_size = tensor.shape[axis]
    if split is None:
        split = 1
        split_indices = [i * split + 1 for i in range(dim_size) if i * split + 1 < dim_size]
        if not keepdims:
            results = np.array_split(tensor, split_indices, axis)
            return [np.squeeze(res, axis) for res in results]
    if np.isscalar(split):
        split_indices = [i * split + 1 for i in range(dim_size) if i * split + 1 < dim_size]  # type: ignore
    else:
        split_indices = np.cumsum(split) + 1
    return np.array_split(tensor, split_indices, axis)  # type: ignore


def ConcatFromSequenceImpl(sequence: List[np.ndarray], axis: int, new_axis: Optional[int] = 0) -> np.ndarray:
    if not new_axis:
        return np.concatenate(sequence, axis)
    else:
        return np.stack(sequence, axis)


class Sequence(Base):

    @staticmethod
    def export() -> None:

        def make_graph(
                nodes: List[onnx.helper.NodeProto],
                input_shapes: List[Optional[typing.Sequence[Union[Text, int]]]],
                output_shapes: List[Optional[typing.Sequence[Union[Text, int]]]],
                input_names: List[Text],
                output_names: List[Text],
                # TODO: Use proper type annotation rather than string
                # once we drop support for Python 3.6.
                # See https://www.python.org/dev/peps/pep-0563/
                input_types: List['TensorProto.DataType'],
                output_types: List['TensorProto.DataType'],
                initializers: Optional[List[TensorProto]] = None
        ) -> onnx.helper.GraphProto:
            graph = onnx.helper.make_graph(
                nodes=nodes,
                name='Sequence',
                inputs=[
                    onnx.helper.make_tensor_value_info(
                        name,
                        input_type,
                        input_shape)
                    for name, input_type, input_shape in zip(input_names, input_types, input_shapes)],
                outputs=[
                    onnx.helper.make_tensor_value_info(
                        name,
                        output_type,
                        output_shape)
                    for name, output_type, output_shape in zip(output_names, output_types, output_shapes)],
                initializer=initializers)
            return graph

        #1st testcase - insert and at.
        # 1. SequenceEmpty:         -> []
        # 2. SequenceInsert(x):     -> [x]
        # 3. SequenceInsert(y):     -> [x, y]
        # 4. SequenceInsert(z, 1):  -> [x, z, y]
        # 5. SequenceAt(2):         -> y
        seq_empty_node = onnx.helper.make_node('SequenceEmpty', [], ['Seq_empty'])
        seq_insert_node = onnx.helper.make_node('SequenceInsert', ['Seq_empty', 'X'], ['Seq_1'])
        seq_insert_node2 = onnx.helper.make_node('SequenceInsert', ['Seq_1', 'Y'], ['Seq_2'])
        seq_insert_node3 = onnx.helper.make_node('SequenceInsert', ['Seq_2', 'Z', 'pos'], ['Seq_3'])
        seq_at_node = onnx.helper.make_node('SequenceAt', ['Seq_3', 'pos_at'], ['out'])

        x_shape = [2, 3, 4]
        y_shape = [1, 3, 4]
        z_shape = [3, 3, 4]
        out_shape = [None, 3, 4]

        x = np.ones(x_shape, dtype=np.float32)
        y = np.zeros(y_shape, dtype=np.float32)
        z = np.ones(z_shape, dtype=np.float32) * 2
        pos_val = 1
        pos_at_val = 2

        out = SequenceEmptyImpl()
        out = SequenceInsertImpl(out, x)
        out = SequenceInsertImpl(out, y)
        out = SequenceInsertImpl(out, z, pos_val)
        out = SequenceAtImpl(out, pos_at_val)
        assert np.array_equal(out, y)

        pos = onnx.helper.make_tensor('pos', TensorProto.INT64, (), (pos_val, ))
        pos_at = onnx.helper.make_tensor('pos_at', TensorProto.INT64, (), (pos_at_val, ))

        graph = make_graph(
            [seq_empty_node, seq_insert_node, seq_insert_node2, seq_insert_node3, seq_at_node],
            [x_shape, y_shape, z_shape, [], []],  # type: ignore
            [out_shape],  # type: ignore
            ['X', 'Y', 'Z', 'pos', 'pos_at'],
            ['out'],
            [onnx.TensorProto.FLOAT] * 3 + [onnx.TensorProto.INT64] * 2,  # type: ignore
            [onnx.TensorProto.FLOAT],
            [pos, pos_at])
        model = onnx.helper.make_model(graph, producer_name='backend-test')
        expect(model, inputs=[x, y, z], outputs=[out], name="test_sequence_model1")

        #2nd testcase - erase and at.
        # 1. SequenceConstruct(x, y, z):    -> [x, y, z]
        # 2. SequenceErase(1):              -> [x, z]
        # 3. SequenceAt(1):                 -> z
        seq_construct_node = onnx.helper.make_node('SequenceConstruct', ['X', 'Y', 'Z'], ['seq_1'])
        seq_erase_node = onnx.helper.make_node('SequenceErase', ['seq_1', 'pos_erase'], ['seq_2'])
        seq_at_node = onnx.helper.make_node('SequenceAt', ['seq_2', 'pos_at'], ['out'])

        tensor_shape = [2, 3, 4]

        x = np.ones(tensor_shape, dtype=np.float32)
        y = np.zeros(tensor_shape, dtype=np.float32)
        z = np.ones(tensor_shape, dtype=np.float32) * 2
        pos_erase_val = 1
        pos_at_val = 1

        out = SequenceConstructImpl(x, y, z)
        out = SequenceEraseImpl(out, pos_erase_val)
        out = SequenceAtImpl(out, pos_at_val)
        assert np.array_equal(out, z)

        pos_erase = onnx.helper.make_tensor('pos_erase', TensorProto.INT64, (), (pos_erase_val, ))
        pos_at = onnx.helper.make_tensor('pos_at', TensorProto.INT64, (), (pos_at_val, ))

        graph = make_graph(
            [seq_construct_node, seq_erase_node, seq_at_node],
            [tensor_shape, tensor_shape, tensor_shape, [], []],  # type: ignore
            [tensor_shape],  # type: ignore
            ['X', 'Y', 'Z', 'pos_erase', 'pos_at'],
            ['out'],
            [onnx.TensorProto.FLOAT] * 3 + [onnx.TensorProto.INT64] * 2,  # type: ignore
            [onnx.TensorProto.FLOAT],
            [pos_erase, pos_at])
        model = onnx.helper.make_model(graph, producer_name='backend-test')
        expect(model, inputs=[x, y, z], outputs=[out], name="test_sequence_model2")

        #3rd testcase - erase, insert and at, with negative index value.
        # 1. SequenceConstruct(x, y, z):    -> [x, y, z]
        # 2. SequenceErase(-3):             -> [y, z]
        # 3. SequenceInsert(x, -1):         -> [y, x, z]
        # 4. SequenceAt(-1):                -> z
        seq_construct_node = onnx.helper.make_node('SequenceConstruct', ['X', 'Y', 'Z'], ['seq_1'])
        seq_erase_node = onnx.helper.make_node('SequenceErase', ['seq_1', 'pos_erase'], ['seq_2'])
        seq_insert_node = onnx.helper.make_node('SequenceInsert', ['seq_2', 'X', 'pos_insert'], ['seq_3'])
        seq_at_node = onnx.helper.make_node('SequenceAt', ['seq_3', 'pos_at'], ['out'])

        tensor_shape = [2, 3, 4]

        x = np.ones(tensor_shape, dtype=np.float32)
        y = np.zeros(tensor_shape, dtype=np.float32)
        z = np.ones(tensor_shape, dtype=np.float32) * 2
        pos_erase_val = -3
        pos_insert_val = -1
        pos_at_val = -1
        out = SequenceConstructImpl(x, y, z)
        out = SequenceEraseImpl(out, pos_erase_val)
        out = SequenceInsertImpl(out, x, pos_insert_val)
        out = SequenceAtImpl(out, pos_at_val)
        assert np.array_equal(out, z)

        pos_erase = onnx.helper.make_tensor('pos_erase', TensorProto.INT64, (), (pos_erase_val, ))
        pos_insert = onnx.helper.make_tensor('pos_insert', TensorProto.INT64, (), (pos_insert_val, ))
        pos_at = onnx.helper.make_tensor('pos_at', TensorProto.INT64, (), (pos_at_val, ))

        graph = make_graph(
            [seq_construct_node, seq_erase_node, seq_insert_node, seq_at_node],
            [tensor_shape, tensor_shape, tensor_shape, [], [], []],  # type: ignore
            [tensor_shape],  # type: ignore
            ['X', 'Y', 'Z', 'pos_erase', 'pos_insert', 'pos_at'],
            ['out'],
            [onnx.TensorProto.FLOAT] * 3 + [onnx.TensorProto.INT64] * 3,  # type: ignore
            [onnx.TensorProto.FLOAT],
            [pos_erase, pos_insert, pos_at])
        model = onnx.helper.make_model(graph, producer_name='backend-test')
        expect(model, inputs=[x, y, z], outputs=[out], name="test_sequence_model3")

        #4th testcase - concat
        seq_construct_node = onnx.helper.make_node('SequenceConstruct', ['X', 'Y', 'Z'], ['seq_1'])
        seq_concat_node = onnx.helper.make_node('ConcatFromSequence', ['seq_1'], ['out'], axis=1)

        tensor_shape = [2, 3, 4]
        concat_out_shape = [2, None, 4]

        x = np.ones(tensor_shape, dtype=np.float32)
        y = np.zeros(tensor_shape, dtype=np.float32)
        z = np.ones(tensor_shape, dtype=np.float32) * 2
        out = SequenceConstructImpl(x, y, z)
        concat_out = ConcatFromSequenceImpl(out, 1)

        graph = make_graph(
            [seq_construct_node, seq_concat_node],
            [tensor_shape] * 3,  # type: ignore
            [concat_out_shape],  # type: ignore
            ['X', 'Y', 'Z'],
            ['out'],
            [onnx.TensorProto.FLOAT] * 3,  # type: ignore
            [onnx.TensorProto.FLOAT])
        model = onnx.helper.make_model(graph, producer_name='backend-test')
        expect(model, inputs=[x, y, z], outputs=[concat_out], name="test_sequence_model4")

        #5th testcase - concat with new_axis = 1
        seq_construct_node = onnx.helper.make_node('SequenceConstruct', ['X', 'Y', 'Z'], ['seq_1'])
        seq_concat_node = onnx.helper.make_node('ConcatFromSequence', ['seq_1'], ['out'], axis=-1, new_axis=1)

        tensor_shape = [2, 3, 4]
        concat_out_shape = [2, 3, 4, 3]

        x = np.ones(tensor_shape, dtype=np.float32)
        y = np.zeros(tensor_shape, dtype=np.float32)
        z = np.ones(tensor_shape, dtype=np.float32) * 2
        out = SequenceConstructImpl(x, y, z)
        concat_out = ConcatFromSequenceImpl(out, -1, 1)

        graph = make_graph(
            [seq_construct_node, seq_concat_node],
            [tensor_shape] * 3,  # type: ignore
            [concat_out_shape],  # type: ignore
            ['X', 'Y', 'Z'],
            ['out'],
            [onnx.TensorProto.FLOAT] * 3,  # type: ignore
            [onnx.TensorProto.FLOAT],)
        model = onnx.helper.make_model(graph, producer_name='backend-test')
        expect(model, inputs=[x, y, z], outputs=[concat_out], name="test_sequence_model5")

        #6th testcase - split and len
        seq_split_node = onnx.helper.make_node('SplitToSequence', ['X'], ['seq_1'], axis=-1)
        seq_len_node = onnx.helper.make_node('SequenceLength', ['seq_1'], ['len'])

        tensor_shape = [2, 3, 4]
        len_shape = []  # type: ignore

        x = np.ones(tensor_shape, dtype=np.float32)
        out = SplitToSequenceImpl(x, axis=-1)
        out = SequenceLengthImpl(out)
        assert np.array_equal(out, np.int64(4))

        graph = onnx.helper.make_graph(
            nodes=[seq_split_node, seq_len_node],
            name='Sequence',
            inputs=[
                onnx.helper.make_tensor_value_info(
                    'X',
                    onnx.TensorProto.FLOAT,
                    tensor_shape)],
            outputs=[
                onnx.helper.make_tensor_value_info(
                    'len',
                    onnx.TensorProto.INT64,
                    len_shape)])  # type: ignore

        model = onnx.helper.make_model(graph, producer_name='backend-test')
        expect(model, inputs=[x], outputs=[out], name="test_sequence_model6")

        #7th testcase - split with keepdims=0, and SequenceAt
        seq_split_node = onnx.helper.make_node('SplitToSequence', ['X'], ['seq_1'], axis=0, keepdims=0)
        seq_at_node = onnx.helper.make_node('SequenceAt', ['seq_1', 'pos_at'], ['out'])

        tensor_shape = [2, 3, 4]
        out_shape = [3, 4]

        x = np.random.rand(*tensor_shape)
        pos_at_val = 1
        out = SplitToSequenceImpl(x, axis=0, keepdims=0)
        out = SequenceAtImpl(out, pos_at_val)
        assert np.array_equal(out, x[pos_at_val])

        pos_at = onnx.helper.make_tensor('pos_at', TensorProto.INT64, (), (pos_at_val, ))

        graph = make_graph(
            [seq_split_node, seq_at_node],
            [tensor_shape, []],  # type: ignore
            [out_shape],  # type: ignore
            ['X', 'pos_at'],
            ['out'],
            [onnx.TensorProto.DOUBLE, onnx.TensorProto.INT64],
            [onnx.TensorProto.DOUBLE],
            [pos_at])
        model = onnx.helper.make_model(graph, producer_name='backend-test')
        expect(model, inputs=[x], outputs=[out], name="test_sequence_model7")

        #8th testcase - split zero length
        seq_split_node = onnx.helper.make_node('SplitToSequence', ['X', 'Splits'], ['seq_1'])
        seq_len_node = onnx.helper.make_node('SequenceLength', ['seq_1'], ['len'])

        tensor_shape = ['n']  # type: ignore
        splits_shape = [3]  # type: ignore

        x = np.array([]).astype(np.float32)
        splits = np.array([0, 0, 0]).astype(np.int64)
        out_len = np.int64(3)

        graph = onnx.helper.make_graph(
            nodes=[seq_split_node, seq_len_node],
            name='Sequence',
            inputs=[
                onnx.helper.make_tensor_value_info(
                    'X',
                    onnx.TensorProto.FLOAT,
                    tensor_shape),  # type: ignore
                onnx.helper.make_tensor_value_info(
                    'Splits',
                    onnx.TensorProto.INT64,
                    splits_shape)],  # type: ignore
            outputs=[
                onnx.helper.make_tensor_value_info(
                    'len',
                    onnx.TensorProto.INT64,
                    len_shape)])  # type: ignore

        model = onnx.helper.make_model(graph, producer_name='backend-test')
        expect(model, inputs=[x, splits], outputs=[out_len], name="test_sequence_model8")
