# 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 .architecture_features import ArchitectureFeatures
from .architecture_features import Block
from .architecture_features import SharedBufferArea
from .architecture_features import SHRAMElements
from .errors import VelaError
from .ethos_u55_regs.ethos_u55_regs import resampling_mode
from .operation import Kernel
from .operation import NpuBlockType
from .range_set import MemoryRangeSet
from .tensor import MemArea


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()

        self.kernel = Kernel(1, 1)
        self.is_elementwise = ps.npu_block_type == NpuBlockType.ElementWise
        self.uses_lut = False
        self.ifm_count = 1

        if ps.primary_op:
            self.kernel = ps.primary_op.kernel
            self.uses_lut = ps.primary_op.activation_lut is not None

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

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

        self.ifm_resampling_mode = resampling_mode.NONE
        self.ifm_bits = 0
        self.ifm_depth = 0
        if ifm_tensor:
            self.ifm_resampling_mode = ifm_tensor.resampling_mode
            self.ifm_bits = ifm_tensor.dtype.size_in_bits()

            if ifm_tensor.shape != []:
                self.ifm_depth = ifm_tensor.shape[-1]

            if self.is_elementwise:
                self.ifm_count = 2
                if ifm_tensor.shape == []:  # Scalar in ifm1
                    assert ifm2_tensor
                    self.ifm_depth = ifm2_tensor.shape[-1]
                    self.ifm_count = 1
                elif not ifm2_tensor or ifm2_tensor.shape == []:  # Scalar in ifm2
                    self.ifm_count = 1

            if self.ifm_bits == 16:
                if is_acc_40bits_used(ps.npu_block_type, ifm_tensor, ofm_tensor, ifm2_tensor):
                    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
                )
            elif self.ifm_bits == 32:
                assert (
                    self.is_elementwise or ps.npu_block_type == NpuBlockType.ReduceSum
                ), "Unsupported 32-bit IFM operation"
                self.use_ifm_element = SHRAMElements.IFM32
            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.available_shram_banks(self.uses_lut) - 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_resampling_mode=self.ifm_resampling_mode
        )
        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

        acc_banks = ofm_config.banks[self.use_accumulator_element]

        # Update bank counts for IFM and Accumulator
        self.banks_required[SharedBufferArea.IFM] = ifm_config.banks[self.use_ifm_element] * self.ifm_count
        self.banks_required[SharedBufferArea.Accumulators] = 0 if self.is_elementwise else acc_banks

        # 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 get_shram_memory_access_range(self):
        # Returns the SHRAM memory access range used by this shared buffer,
        # excluding access to LUT
        return MemoryRangeSet(
            MemArea.Shram, 0, self.arch.available_shram_banks(self.uses_lut) * self.arch.shram_bank_size
        )


def is_acc_40bits_used(npu_block_type, ifm_tensor, ofm_tensor, ifm2_tensor=None):
    tensors = [t for t in (ifm_tensor, ifm2_tensor, ofm_tensor) if t is not None]
    scales = [t.quantization.scale_f32 for t in tensors if t.quantization is not None]
    has_scale = len(tensors) == len(scales) and None not in scales
    return npu_block_type != NpuBlockType.Pooling and has_scale


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)
        if config is None:
            raise VelaError("Block config override '{0}' cannot be allocated".format(arch.override_block_config))
        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