from sympy.tensor.functions import TensorProduct
from sympy import MatrixSymbol, Matrix, Array
from sympy.abc import x, y, z
from sympy.abc import i, j, k, l


A = MatrixSymbol("A", 3, 3)
B = MatrixSymbol("B", 3, 3)
C = MatrixSymbol("C", 3, 3)


def test_TensorProduct_construction():
    assert TensorProduct(3, 4) == 12
    assert isinstance(TensorProduct(A, A), TensorProduct)

    expr = TensorProduct(TensorProduct(x, y), z)
    assert expr == x*y*z

    expr = TensorProduct(TensorProduct(A, B), C)
    assert expr == TensorProduct(A, B, C)

    expr = TensorProduct(Matrix.eye(2), [[0, -1], [1, 0]])
    assert expr == Array([
        [
            [[0, -1], [1, 0]],
            [[0, 0], [0, 0]]
        ],
        [
            [[0, 0], [0, 0]],
            [[0, -1], [1, 0]]
        ]
    ])


def test_TensorProduct_shape():

    expr = TensorProduct(3, 4, evaluate=False)
    assert expr.shape == ()
    assert expr.rank() == 0

    expr = TensorProduct([1, 2], [x, y], evaluate=False)
    assert expr.shape == (2, 2)
    assert expr.rank() == 2
    expr = TensorProduct(expr, expr, evaluate=False)
    assert expr.shape == (2, 2, 2, 2)
    assert expr.rank() == 4

    expr = TensorProduct(Matrix.eye(2), [[0, -1], [1, 0]], evaluate=False)
    assert expr.shape == (2, 2, 2, 2)
    assert expr.rank() == 4


def test_TensorProduct_getitem():
    expr = TensorProduct(A, B)
    assert expr[i, j, k, l] == A[i, j]*B[k, l]

[((7, 4, 7, 27), 'sympy.MatrixSymbol', 'MatrixSymbol', ({(7, 17, 7, 20): '"""A"""', (7, 22, 7, 23): '3', (7, 25, 7, 26): '3'}, {}), "('A', 3, 3)", False, 'from sympy import MatrixSymbol, Matrix, Array\n'), ((8, 4, 8, 27), 'sympy.MatrixSymbol', 'MatrixSymbol', ({(8, 17, 8, 20): '"""B"""', (8, 22, 8, 23): '3', (8, 25, 8, 26): '3'}, {}), "('B', 3, 3)", False, 'from sympy import MatrixSymbol, Matrix, Array\n'), ((9, 4, 9, 27), 'sympy.MatrixSymbol', 'MatrixSymbol', ({(9, 17, 9, 20): '"""C"""', (9, 22, 9, 23): '3', (9, 25, 9, 26): '3'}, {}), "('C', 3, 3)", False, 'from sympy import MatrixSymbol, Matrix, Array\n'), ((37, 11, 37, 46), 'sympy.tensor.functions.TensorProduct', 'TensorProduct', (), '', False, 'from sympy.tensor.functions import TensorProduct\n'), ((41, 11, 41, 56), 'sympy.tensor.functions.TensorProduct', 'TensorProduct', (), '', False, 'from sympy.tensor.functions import TensorProduct\n'), ((44, 11, 44, 52), 'sympy.tensor.functions.TensorProduct', 'TensorProduct', (), '', False, 'from sympy.tensor.functions import TensorProduct\n'), ((54, 11, 54, 30), 'sympy.tensor.functions.TensorProduct', 'TensorProduct', ({(54, 25, 54, 26): 'A', (54, 28, 54, 29): 'B'}, {}), '(A, B)', False, 'from sympy.tensor.functions import TensorProduct\n'), ((13, 11, 13, 30), 'sympy.tensor.functions.TensorProduct', 'TensorProduct', ({(13, 25, 13, 26): '(3)', (13, 28, 13, 29): '(4)'}, {}), '(3, 4)', False, 'from sympy.tensor.functions import TensorProduct\n'), ((14, 22, 14, 41), 'sympy.tensor.functions.TensorProduct', 'TensorProduct', ({(14, 36, 14, 37): 'A', (14, 39, 14, 40): 'A'}, {}), '(A, A)', False, 'from sympy.tensor.functions import TensorProduct\n'), ((16, 25, 16, 44), 'sympy.tensor.functions.TensorProduct', 'TensorProduct', ({(16, 39, 16, 40): 'x', (16, 42, 16, 43): 'y'}, {}), '(x, y)', False, 'from sympy.tensor.functions import TensorProduct\n'), ((19, 25, 19, 44), 'sympy.tensor.functions.TensorProduct', 'TensorProduct', ({(19, 39, 19, 40): 'A', (19, 42, 19, 43): 'B'}, {}), '(A, B)', False, 'from sympy.tensor.functions import TensorProduct\n'), ((20, 19, 20, 41), 'sympy.tensor.functions.TensorProduct', 'TensorProduct', ({(20, 33, 20, 34): 'A', (20, 36, 20, 37): 'B', (20, 39, 20, 40): 'C'}, {}), '(A, B, C)', False, 'from sympy.tensor.functions import TensorProduct\n'), ((22, 25, 22, 38), 'sympy.Matrix.eye', 'Matrix.eye', ({(22, 36, 22, 37): '2'}, {}), '(2)', False, 'from sympy import MatrixSymbol, Matrix, Array\n'), ((23, 19, 32, 6), 'sympy.Array', 'Array', ({(23, 25, 32, 5): '[[[[0, -1], [1, 0]], [[0, 0], [0, 0]]], [[[0, 0], [0, 0]], [[0, -1], [1, 0]]]]'}, {}), '([[[[0, -1], [1, 0]], [[0, 0], [0, 0]]], [[[0, 0], [0, 0]], [[0, -1],\n    [1, 0]]]])', False, 'from sympy import MatrixSymbol, Matrix, Array\n'), ((48, 25, 48, 38), 'sympy.Matrix.eye', 'Matrix.eye', ({(48, 36, 48, 37): '2'}, {}), '(2)', False, 'from sympy import MatrixSymbol, Matrix, Array\n')]