# Copyright (C) 2018-2021 Intel Corporation
# SPDX-License-Identifier: Apache-2.0

import unittest

import numpy as np
from generator import generator, generate

from extensions.ops.einsum import Einsum
from mo.front.common.partial_infer.utils import int64_array
from mo.graph.graph import Graph
from mo.graph.graph import Node
from unit_tests.utils.graph import build_graph, regular_op_with_shaped_data, result, connect


def create_einsum_graph(input_shapes: list, equation: str) -> Graph:
    num_inputs = len(input_shapes)
    assert num_inputs > 0, "Einsum node must have at least one input"
    nodes = {}
    edges = []
    for input_ind in range(num_inputs):
        input_name = 'input' + str(input_ind)
        parameter_op = regular_op_with_shaped_data(input_name, input_shapes[input_ind],
                                                   {'op': 'Parameter', 'type': 'Parameter'})
        nodes.update(parameter_op)
        edges += connect(input_name, str(input_ind) + ":einsum_node")
    einsum_op = regular_op_with_shaped_data('einsum_node', None,
                                            {'op': 'Einsum', 'type': 'Einsum', 'equation': equation})
    nodes.update(einsum_op)
    result_op = result('output')
    nodes.update(result_op)
    edges += connect('einsum_node', 'output')

    graph = build_graph(nodes, edges, nodes_with_edges_only=True)
    return graph


@generator
class TestEinsum(unittest.TestCase):
    @generate(*[
        # dot product
        ([int64_array([10]), int64_array([10])], "i,i->", int64_array([])),
        # matrix multiplication
        ([int64_array([2, 3]), int64_array([3, 4])], "ab,bc->ac", int64_array([2, 4])),
        # trace per batch
        ([int64_array([2, 3, 3])], "kii->k", int64_array([2])),
        # diagonal extraction
        ([int64_array([6, 5, 5])], "kii->ki", int64_array([6, 5])),
        # transpose
        ([int64_array([1, 2, 3])], "ijk->kij", int64_array([3, 1, 2])),
        # multiple matrix multiplication
        ([int64_array([2, 5]), int64_array([5, 3, 6]), int64_array([5, 3])], "ab,bcd,bc->ca", int64_array([3, 2])),
        # ellipsis for one operand
        ([int64_array([5, 3, 4])], "a...->...", int64_array([3, 4])),
        # ellipsis for multiple operands
        ([int64_array([3, 5]), int64_array([1])], "a...,...->a...", int64_array([3, 5])),
        # ellipsis with broadcasting
        ([int64_array([9, 1, 4, 3]), int64_array([3, 11, 7, 1])], "a...b,b...->a...", int64_array([9, 11, 7, 4])),
        # mixed case letters in equation
        ([int64_array([1, 3, 5])], "AbC", int64_array([1, 5, 3])),
        # mixed case letters and equation in implicit mode
        ([int64_array([3, 11, 1, 5]), int64_array([1, 3, 1, 7])], "a...b,B...", int64_array([3, 11, 7, 1, 3, 5])),
        # inner product in implicit mode
        ([int64_array([3]), int64_array([3])], "i,i", int64_array([])),
        # equation with ellipsis and repeated labels in implicit mode
        # "a...b,b..." is equivalent to "a...b,b...->...a"
        ([int64_array([9, 1, 4, 3]), int64_array([3, 11, 7, 1])], "a...b,b...", int64_array([11, 7, 4, 9])),
    ])
    def test_einsum(self, input_shapes, equation, ref_output_shape):
        graph = create_einsum_graph(input_shapes, equation)
        einsum_node = Node(graph, 'einsum_node')
        Einsum.infer(einsum_node)

        # get the result
        res_output_shape = graph.node['einsum_node_d']['shape']

        self.assertTrue(np.array_equal(ref_output_shape, res_output_shape),
                        'shape does not match expected: {} and given: {}'.format(ref_output_shape, res_output_shape))

    @generate(*[
        # incorrect subscript numbers or inputs
        ([int64_array([3, 11]), int64_array([11, 4])], "ab,bc,cd->ac", None),
        # invalid labels
        ([int64_array([3, 11]), int64_array([11, 4])], "a$,Bc->ac", None),
        # incompatible shapes
        ([int64_array([3, 11]), int64_array([12, 4])], "ab,bc->ac", None),
        # not broadcastable shapes
        ([int64_array([11, 1, 4, 3]), int64_array([3, 11, 7, 5])], "a...b,b...->a...", None),
        # missed ellipsis
        ([int64_array([11, 1, 4, 3]), int64_array([3, 11, 7, 4])], "a...b,b...->a", None),
    ])
    def test_invalid_cases(self, input_shapes, equation, ref_output_shape):
        graph = create_einsum_graph(input_shapes, equation)
        einsum_node = Node(graph, 'einsum_node')
        self.assertRaises(AssertionError, Einsum.infer, einsum_node)
