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]

[((164, 187), 'sympy.MatrixSymbol', 'MatrixSymbol', (['"""A"""', '(3)', '(3)'], {}), "('A', 3, 3)\n", (176, 187), False, 'from sympy import MatrixSymbol, Matrix, Array\n'), ((192, 215), 'sympy.MatrixSymbol', 'MatrixSymbol', (['"""B"""', '(3)', '(3)'], {}), "('B', 3, 3)\n", (204, 215), False, 'from sympy import MatrixSymbol, Matrix, Array\n'), ((220, 243), 'sympy.MatrixSymbol', 'MatrixSymbol', (['"""C"""', '(3)', '(3)'], {}), "('C', 3, 3)\n", (232, 243), False, 'from sympy import MatrixSymbol, Matrix, Array\n'), ((848, 883), 'sympy.tensor.functions.TensorProduct', 'TensorProduct', (['(3)', '(4)'], {'evaluate': '(False)'}), '(3, 4, evaluate=False)\n', (861, 883), False, 'from sympy.tensor.functions import TensorProduct\n'), ((952, 997), 'sympy.tensor.functions.TensorProduct', 'TensorProduct', (['[1, 2]', '[x, y]'], {'evaluate': '(False)'}), '([1, 2], [x, y], evaluate=False)\n', (965, 997), False, 'from sympy.tensor.functions import TensorProduct\n'), ((1069, 1110), 'sympy.tensor.functions.TensorProduct', 'TensorProduct', (['expr', 'expr'], {'evaluate': '(False)'}), '(expr, expr, evaluate=False)\n', (1082, 1110), False, 'from sympy.tensor.functions import TensorProduct\n'), ((1366, 1385), 'sympy.tensor.functions.TensorProduct', 'TensorProduct', (['A', 'B'], {}), '(A, B)\n', (1379, 1385), False, 'from sympy.tensor.functions import TensorProduct\n'), ((296, 315), 'sympy.tensor.functions.TensorProduct', 'TensorProduct', (['(3)', '(4)'], {}), '(3, 4)\n', (309, 315), False, 'from sympy.tensor.functions import TensorProduct\n'), ((344, 363), 'sympy.tensor.functions.TensorProduct', 'TensorProduct', (['A', 'A'], {}), '(A, A)\n', (357, 363), False, 'from sympy.tensor.functions import TensorProduct\n'), ((406, 425), 'sympy.tensor.functions.TensorProduct', 'TensorProduct', (['x', 'y'], {}), '(x, y)\n', (419, 425), False, 'from sympy.tensor.functions import TensorProduct\n'), ((481, 500), 'sympy.tensor.functions.TensorProduct', 'TensorProduct', (['A', 'B'], {}), '(A, B)\n', (494, 500), False, 'from sympy.tensor.functions import TensorProduct\n'), ((524, 546), 'sympy.tensor.functions.TensorProduct', 'TensorProduct', (['A', 'B', 'C'], {}), '(A, B, C)\n', (537, 546), False, 'from sympy.tensor.functions import TensorProduct\n'), ((573, 586), 'sympy.Matrix.eye', 'Matrix.eye', (['(2)'], {}), '(2)\n', (583, 586), False, 'from sympy import MatrixSymbol, Matrix, Array\n'), ((626, 715), 'sympy.Array', 'Array', (['[[[[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]]]])\n', (631, 715), False, 'from sympy import MatrixSymbol, Matrix, Array\n'), ((1203, 1216), 'sympy.Matrix.eye', 'Matrix.eye', (['(2)'], {}), '(2)\n', (1213, 1216), False, 'from sympy import MatrixSymbol, Matrix, Array\n')]