# Copyright (C) 2020-2021 Arm Limited or its affiliates. All rights reserved.
#
# SPDX-License-Identifier: Apache-2.0
#
# Licensed 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
#
# 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.
# Description:
# Defines the class Shape4D.
from collections import namedtuple

from .numeric_util import full_shape
from .numeric_util import round_up_divide


class Shape4D(namedtuple("Shape4D", ["batch", "height", "width", "depth"])):
    """
    4D Shape (in NHWC format)
    """

    def __new__(cls, n=1, h=1, w=1, c=1):
        assert n is not None
        if isinstance(n, list):
            assert h == 1 and w == 1 and c == 1
            tmp = full_shape(4, n, 1)
            self = super(Shape4D, cls).__new__(cls, tmp[0], tmp[1], tmp[2], tmp[3])
        else:
            self = super(Shape4D, cls).__new__(cls, n, h, w, c)
        return self

    @classmethod
    def from_list(cls, shape, base=1):
        tmp = full_shape(4, shape, base)
        return cls(tmp[0], tmp[1], tmp[2], tmp[3])

    @classmethod
    def from_hwc(cls, h, w, c):
        return cls(1, h, w, c)

    def with_batch(self, new_batch):
        return Shape4D(new_batch, self.height, self.width, self.depth)

    def with_height(self, new_height):
        return Shape4D(self.batch, new_height, self.width, self.depth)

    def with_width(self, new_width):
        return Shape4D(self.batch, self.height, new_width, self.depth)

    def with_hw(self, new_height, new_width):
        return Shape4D(self.batch, new_height, new_width, self.depth)

    def with_depth(self, new_depth):
        return Shape4D(self.batch, self.height, self.width, new_depth)

    def add(self, n, h, w, c):
        return Shape4D(self.batch + n, self.height + h, self.width + w, self.depth + c)

    def __add__(self, rhs):
        return Shape4D(self.batch + rhs.batch, self.height + rhs.height, self.width + rhs.width, self.depth + rhs.depth)

    def __sub__(self, rhs):
        return Shape4D(self.batch - rhs.batch, self.height - rhs.height, self.width - rhs.width, self.depth - rhs.depth)

    def __floordiv__(self, rhs):
        return Shape4D(
            self.batch // rhs.batch, self.height // rhs.height, self.width // rhs.width, self.depth // rhs.depth
        )

    def __mod__(self, rhs):
        return Shape4D(self.batch % rhs.batch, self.height % rhs.height, self.width % rhs.width, self.depth % rhs.depth)

    def __str__(self):
        return f"<Shape4D {list(self)}>"

    def div_round_up(self, rhs):
        return Shape4D(
            round_up_divide(self.batch, rhs.batch),
            round_up_divide(self.height, rhs.height),
            round_up_divide(self.width, rhs.width),
            round_up_divide(self.depth, rhs.depth),
        )

    def elements(self):
        return self.batch * self.width * self.height * self.depth

    def elements_wh(self):
        return self.width * self.height

    def is_empty(self):
        return (self.batch + self.width + self.height + self.depth) == 0

    def as_list(self):
        return list(self)

    def get_hw_as_list(self):
        return list([self.height, self.width])