# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=invalid-name, line-too-long, unused-variable, too-many-locals
"""Reorg in python"""

import numpy as np


def reorg_python(a_np, stride):
    """Reorg operator

    Parameters
    ----------
    a_np : numpy.ndarray
        4-D with shape [batch, in_channel, in_height, in_width]

    stride : int
        Stride size

    Returns
    -------
    b_np : np.ndarray
        4-D with shape [batch, out_channel, out_height, out_width]
    """

    batch, in_channel, in_height, in_width = a_np.shape
    a_np = np.reshape(a_np, batch * in_channel * in_height * in_width)
    out_c = int(in_channel / (stride * stride))
    out_channel = in_channel * stride * stride
    out_height = int(in_height / stride)
    out_width = int(in_width / stride)
    b_np = np.zeros(batch * out_channel * out_height * out_width)
    cnt = 0
    for b in range(batch):
        for k in range(in_channel):
            for j in range(in_height):
                for i in range(in_width):
                    c2 = k % out_c
                    offset = int(k / out_c)
                    w2 = int(i * stride + offset % stride)
                    h2 = int(j * stride + offset / stride)
                    out_index = int(
                        w2 + in_width * stride * (h2 + in_height * stride * (c2 + out_c * b))
                    )
                    b_np[cnt] = a_np[int(out_index)]
                    cnt = cnt + 1
    b_np = np.reshape(b_np, (batch, out_channel, out_height, out_width))
    return b_np