# Copyright (C) 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:
# Unit tests for architecture_allocator.py
import pytest

from ethosu.vela.architecture_allocator import find_block_config
from ethosu.vela.architecture_allocator import try_block_config
from ethosu.vela.architecture_features import Accelerator
from ethosu.vela.architecture_features import Block
from ethosu.vela.architecture_features import create_default_arch
from ethosu.vela.ethos_u55_regs.ethos_u55_regs import resampling_mode
from ethosu.vela.operation import Kernel
from ethosu.vela.operation import NpuBlockType
from ethosu.vela.shape4d import Shape4D

test_data = [
    {
        "block_type": NpuBlockType.ConvolutionDepthWise,
        "kernel": Kernel(25, 5, 2, 2, 1, 1),
        "ofm_shape": Shape4D(2, 11, 22),
        "ifm_shape": Shape4D(27, 25, 22),
    },
    {
        "block_type": NpuBlockType.Pooling,
        "kernel": Kernel(2, 2),
        "ofm_shape": Shape4D(53, 49, 22),
        "ifm_shape": Shape4D(27, 25, 22),
        "ifm_resampling": resampling_mode.NEAREST,
    },
    {
        "block_type": NpuBlockType.ConvolutionMxN,
        "accelerator": Accelerator.Ethos_U55_32,
        "kernel": Kernel(2, 5),
        "ofm_shape": Shape4D(48, 1, 17),
        "ifm_shape": Shape4D(24, 5, 18),
        "ifm_resampling": resampling_mode.TRANSPOSE,
    },
    {
        "block_type": NpuBlockType.ElementWise,
        "ofm_shape": Shape4D(27, 2, 22),
        "ifm_shape": Shape4D(27, 2, 1),
        "ifm2_shape": Shape4D(27, 25, 22),
    },
    {
        "block_type": NpuBlockType.ElementWise,
        "accelerator": Accelerator.Ethos_U55_32,
        "ofm_shape": Shape4D(48, 37, 17),
        "ifm_shape": Shape4D(48, 37, 17),
        "uses_scalar": True,
        "lut_banks": 2,
    },
    {
        "block_type": NpuBlockType.ElementWise,
        "ofm_shape": Shape4D(27, 2, 22),
        "ifm_shape": Shape4D(27, 2, 22),
        "ifm_bits": 16,
    },
]


@pytest.mark.parametrize("test_data", test_data)
def test_allocate(test_data):
    """Tests that find_block_config and try_block_config produce consistent SHRAM layouts"""
    accelerator = test_data.get("accelerator", Accelerator.Ethos_U55_128)
    arch = create_default_arch(accelerator)
    kernel = test_data.get("kernel", Kernel(1, 1))
    block_type = test_data["block_type"]
    ofm_shape = test_data["ofm_shape"]
    ifm_shape = test_data["ifm_shape"]
    ifm2_shape = test_data.get("ifm2_shape")
    uses_scalar = test_data.get("uses_scalar", False)
    ifm_bits = test_data.get("ifm_bits", 8)
    ifm_resampling = test_data.get("ifm_resampling", resampling_mode.NONE)
    scaled = test_data.get("scaled", True)
    lut_banks = test_data.get("lut_banks", 0)
    config = find_block_config(
        arch,
        block_type,
        ofm_shape,
        ifm_shape,
        ifm2_shape,
        uses_scalar=uses_scalar,
        ifm_bits=ifm_bits,
        kernel=kernel,
        lut_banks=lut_banks,
        scaled=scaled,
        ifm_resampling=ifm_resampling,
    )
    assert config is not None
    config2 = try_block_config(
        Block.from_shape(config.ofm_block.as_list()),
        arch,
        block_type,
        ofm_shape,
        ifm_shape,
        ifm2_shape,
        is_partkernel=config.is_partkernel,
        uses_scalar=uses_scalar,
        ifm_bits=ifm_bits,
        kernel=kernel,
        lut_banks=lut_banks,
        scaled=scaled,
        ifm_resampling=ifm_resampling,
    )
    assert config2 is not None
    assert config.layout.ib_end == config2.layout.ib_end
    assert config.layout.ab_start == config2.layout.ab_start
    assert config.layout.ib_start2 == config2.layout.ib_start2
    assert config.acc_type == config2.acc_type