# Copyright (C) 2020 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:
# Shared buffer allocation works out how to allocate the Ethos-U55 shared buffer for a given pass.

import numpy as np
from .nn_graph import NpuBlockType
from .numeric_util import round_up_divide, round_up
from .architecture_features import Block, Kernel, SHRAMElements, SharedBufferArea, ArchitectureFeatures
from . import pass_packing


class SharedBufferAllocation:
    def __init__(self, arch, ps):
        self.arch = arch

        self.bank_locations = np.zeros(SharedBufferArea.Size)
        self.banks_required = np.zeros(SharedBufferArea.Size)

        ifm_tensor, ifm2_tensor, weight_tensor, ofm_tensor = ps.get_primary_op_ifm_ifm2_weights_ofm()

        strides = (1, 1, 1, 1)
        dilation = (1, 1, 1, 1)
        self.kernel = Kernel(1, 1)
        is_elementwise = ps.npu_block_type == NpuBlockType.ElementWise

        if ps.primary_op:
            strides = ps.primary_op.attrs.get("strides", strides)
            dilation = ps.primary_op.attrs.get("dilation", dilation)
            k_h = 1
            k_w = 1
            if weight_tensor:
                if ps.primary_op.type != "FullyConnectedAct":
                    k_h = weight_tensor.shape[0]
                    k_w = weight_tensor.shape[1]
            else:
                k_h = ps.primary_op.attrs.get("filter_height", 1)
                k_w = ps.primary_op.attrs.get("filter_width", 1)

            self.kernel = Kernel(k_w, k_h, strides[2], strides[1], dilation[2], dilation[1])

        self.is_equal_depth_op = is_elementwise or ps.npu_block_type in (
            NpuBlockType.ConvolutionDepthWise,
            NpuBlockType.Pooling,
        )
        self.strides = strides

        self.use_accumulator_element = SHRAMElements.Acc32
        if is_elementwise:
            self.use_ifm_element = SHRAMElements.IFM8_Elementwise
        else:
            self.use_ifm_element = SHRAMElements.IFM8

        self.ifm_bits = 0
        self.ifm_depth = 0
        if ifm_tensor:
            self.ifm_bits = ifm_tensor.dtype.size_in_bits()
            if ifm_tensor.shape == [] and is_elementwise:
                # Elementwise operator with scalar in ifm, use ifm2 depth
                self.ifm_depth = ifm2_tensor.shape[-1]
            else:
                self.ifm_depth = ifm_tensor.shape[-1]
            if self.ifm_bits == 16:
                self.use_accumulator_element = SHRAMElements.Acc40
                self.use_ifm_element = self.use_ifm_element + 1
                assert (self.use_ifm_element == SHRAMElements.IFM16) or (
                    self.use_ifm_element == SHRAMElements.IFM16_Elementwise
                )
            else:
                assert self.ifm_bits == 8, "Unexpected IFM bitdepth"

        self.ifm_block_depth = arch.calc_ifm_block_depth(self.ifm_depth, self.ifm_bits)
        self.ofm_tensor = ofm_tensor

        self.banks_required[SharedBufferArea.Weights] = arch.shram_reserved_weight_banks
        self.banks_required[SharedBufferArea.OFM] = arch.shram_reserved_output_banks

    def is_valid(self):
        # Assign zero-based bank starts (first element remains zero)
        self.bank_locations[1:] = np.cumsum(self.banks_required)[:-1]

        # Accumulator area is measured from the end of the buffer
        self.bank_locations[SharedBufferArea.Accumulators] = (
            self.arch.shram_total_banks - self.banks_required[SharedBufferArea.Accumulators]
        )
        ifm_end = self.bank_locations[SharedBufferArea.IFM] + self.banks_required[SharedBufferArea.IFM]
        return ifm_end <= self.bank_locations[SharedBufferArea.Accumulators]

    def try_block(self, ofm_block: Block):
        # Get IFM block configuration
        ifm_block_depth = ofm_block.depth if self.is_equal_depth_op else self.ifm_block_depth
        ifm_block = self.arch.get_ifm_block_size(ifm_block_depth, ofm_block, self.kernel)
        ifm_config = self.arch.get_block_config(ifm_block.width, ifm_block.height, ifm_block.depth)
        if ifm_config is None:
            return None

        # Get OFM block configuration
        ofm_config = self.arch.get_block_config(ofm_block.width, ofm_block.height, ofm_block.depth)
        if ofm_config is None:
            return None

        # Update bank counts for IFM and Accumulator
        self.banks_required[SharedBufferArea.IFM] = ifm_config.banks[self.use_ifm_element]
        self.banks_required[SharedBufferArea.Accumulators] = ofm_config.banks[self.use_accumulator_element]

        # Validating calculates bank layout and returns validity
        if not self.is_valid():
            return None

        return (ofm_block.height, ofm_block.width, ifm_block.depth, ofm_block.depth)

    def generate_used_mask(self, active_set):
        res = np.zeros(self.arch.shram_total_banks, dtype=np.int64)
        for kind in active_set:
            start = int(self.bank_locations[kind])
            end = start + int(self.banks_required[kind])
            res[start:end] = 1
        return res

    def is_compatible(first, second):
        """See if the bank allocations of two convolutions are compatible,
        so that they can run back-to-back without a fence in between"""

        first_set = set((SharedBufferArea.OFM, SharedBufferArea.Accumulators))
        second_set = set((SharedBufferArea.IFM, SharedBufferArea.Weights))

        first_mask = first.generate_used_mask(first_set)
        second_mask = second.generate_used_mask(second_set)

        if np.sum(first_mask & second_mask):
            # overlap
            return False

        return True


def shared_buffer_allocation_for_pass_and_block_config(arch, ps, block_config):
    alloc = SharedBufferAllocation(arch, ps)
    assert (alloc.ifm_block_depth == block_config[2]) or alloc.is_equal_depth_op
    if alloc.try_block(Block(block_config[1], block_config[0], block_config[3])):
        return alloc

    return None


def find_block_configs_suitable_for_pass_and_shared_buffer(arch, ps):
    alloc = SharedBufferAllocation(arch, ps)

    if arch.override_block_config:
        config = alloc.try_block(arch.override_block_config)
        assert config, "Block config override cannot be used"
        return [config]

    # Constrain the search space if the OFM is smaller than the max block size
    # - Add other block search constraints here if required
    if len(alloc.ofm_tensor.shape) == 2:
        max_block_height = max_block_width = alloc.ofm_tensor.shape[0]
    else:
        max_block_width = alloc.ofm_tensor.shape[-2]
        max_block_height = alloc.ofm_tensor.shape[-3]

    # Common block depth
    max_block_depth = alloc.ofm_tensor.shape[-1]

    # Constrain to valid ranges before search
    max_block_width = min(arch.ofm_block_max.width, max_block_width)
    max_block_height = min(arch.ofm_block_max.height, max_block_height)
    max_block_depth = min(arch.ofm_block_max.depth, max_block_depth)

    valid_block_configs = []
    # Try a range of block shapes against this pass
    for w in range(arch.ofm_ublock.width, max_block_width + arch.ofm_ublock.width, arch.ofm_ublock.width):
        for h in range(arch.ofm_ublock.height, max_block_height + arch.ofm_ublock.height, arch.ofm_ublock.height):
            # Try valid OFM block depths
            for c in range(arch.ofm_ublock.depth, max_block_depth + arch.ofm_ublock.depth, arch.ofm_ublock.depth):
                # OFM block depth has the constraint that if it causes the OFM to be
                # split, it must be a multiple of the OFM split size
                if (c >= max_block_depth) or (c < max_block_depth and (c % ArchitectureFeatures.OFMSplitDepth) == 0):
                    config = alloc.try_block(Block(w, h, c))
                    if config:
                        valid_block_configs.append(config)

    assert len(valid_block_configs) > 0
    return valid_block_configs