import mindspore.nn
import torch
import numpy as np
import mindspore as ms


def torch_output_fix(input_data, target_shape):
    r"""
    :param input_data:
    :param target_shape:
    :return: input_data in target shape
    example:
    data = torch.tensor(np.random.randn(1, 3, 32, 32))
    output_data = output_fix(data, target_shape=[1, 64, 224, 224])
    """

    # print(input_data.shape)
    # print(target_shape)

    flatten = torch.nn.Flatten()
    elements = get_element_nums(input_data)
    b1 = input_data.shape[0]
    s1 = elements // b1
    target_elements = 1
    for shape in target_shape:
        target_elements *= shape
    b2 = target_shape[0]
    s2 = target_elements // b2
    det = elements - target_elements
    if det == 0:
        return torch.reshape(input_data, target_shape)

    x1 = s2 - s1
    x2 = b2 - b1
    out = flatten(input_data)
    if x1 > 0:
        out = torch.nn.functional.pad(out, (0, x1))
    elif x1 < 0:
        out = torch.index_select(out, dim=-1,
                                 index=torch.tensor([i for i in range(s2)], dtype=torch.int64))

    if x2 > 0:
        out = torch.transpose(out, dim0=1, dim1=0)
        out = torch.nn.functional.pad(out, (0, x2))
        out = torch.transpose(out, dim1=1, dim0=0)
    elif x2 < 0:
        out = torch.index_select(out, dim=0,
                                 index=torch.tensor([i for i in range(b2)], dtype=torch.int64))
    return torch.reshape(out, target_shape)


def ms_output_fix(input_data, target_shape):
    flatten = mindspore.nn.Flatten()
    elements = get_element_nums(input_data)
    b1 = input_data.shape[0]
    s1 = elements // b1
    target_elements = 1
    for shape in target_shape:
        target_elements *= shape
    b2 = target_shape[0]
    s2 = target_elements // b2
    det = elements - target_elements
    x1 = s2 - s1
    x2 = b2 - b1
    if det == 0:
        return ms.ops.Reshape()(input_data, tuple(target_shape))

    out = flatten(input_data)
    if x1 > 0:
        out = ms.nn.Pad(paddings=((0, 0), (0, x1)))(out)
    elif x1 < 0:
        out = ms.ops.Slice()(out, (0, 0), (-1, s2))

    if x2 > 0:
        out = ms.nn.Pad(paddings=((x2, 0), (0, 0)))(out)
    elif x2 < 0:
        out = ms.ops.Slice()(out, (0, 0), (b2, -1))
    return ms.ops.Reshape()(out, tuple(target_shape))


def get_element_nums(data):
    if isinstance(data, tuple):
        data = list(tuple)
    shape = data.shape
    res = 1
    for i in shape:
        res *= i
    return res


if __name__ == '__main__':
    input_size,output_size=(300,400),(200, 210, 500,2,5)

    torch_data = torch.tensor(np.random.randn(*input_size))
    print(torch_output_fix(torch_data, target_shape=output_size).shape)

    ms_data = ms.Tensor(np.random.randn(*input_size))
    print(ms_output_fix(ms_data, target_shape=output_size).shape)