# -*- coding: utf-8 -*-
# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2020 Megvii Inc. All rights reserved.
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
import numpy as np
import pytest

from megengine import tensor
from megengine.test import assertTensorClose


def test_reshape_tuple():
    inp = tensor(np.arange(1, 17, dtype=np.int32).reshape(4, 4))
    out = tensor(np.arange(100, 116, dtype=np.int32).reshape(1, 16))
    out = out.reshape(inp.shape)

    assertTensorClose(out.numpy(), np.arange(100, 116, dtype=np.int32).reshape(4, 4))


def test_reshape_asterisk():
    inp = tensor(np.arange(1, 17, dtype=np.int32).reshape(4, 4))
    out = tensor(np.arange(100, 116, dtype=np.int32).reshape(1, 16))
    out = out.reshape(*inp.shape)

    assertTensorClose(out.numpy(), np.arange(100, 116, dtype=np.int32).reshape(4, 4))


def test_reshape_shapeof():
    inp = tensor(np.arange(1, 17, dtype=np.int32).reshape(4, 4))
    out = tensor(np.arange(100, 116, dtype=np.int32).reshape(1, 16))
    out = out.reshape(inp.shapeof())

    assertTensorClose(out.numpy(), np.arange(100, 116, dtype=np.int32).reshape(4, 4))


def test_reshape_tensor():
    out = tensor(np.arange(100, 116, dtype=np.int32).reshape(1, 16))
    out = out.reshape(tensor([4, 4]))

    assertTensorClose(out.numpy(), np.arange(100, 116, dtype=np.int32).reshape(4, 4))


def test_reshape_tensor_fused():
    out = tensor(np.arange(100, 116, dtype=np.int32).reshape(1, 16))
    out = out.reshape(
tensor([4, 4])