# -------------------------------------------------------------------------
#  This file is part of the MindStudio project.
# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
#
# MindStudio is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#
#          http://license.coscl.org.cn/MulanPSL2
#
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.
# -------------------------------------------------------------------------

from typing import Union, List

from atk.common.log import Logger
from atk.configs.design_config import DesignConfig
from atk.configs.case_config import CaseConfig, InputCaseConfig
from atk.case_generator.generator.parameter_types import (
    ParameterFactory,
    PARAMETER_REGISTRY
)
from atk.case_generator.generator.processor import ShapeProcessor, NumberProcessor
from atk.common.utils import cal_tensor_numel

logging = Logger().get_logger()
MAX_SHAPE_LENGTH = 2 ** 32 - 1
MAX_SHAPE_LENGTH_FOR_TENSORS = 2 ** 32 - 1


class CaseGenerator:
    def __init__(self, config: DesignConfig):
        self.config = config
        self.index = 0
        self.is_gen_extra = getattr(self.config, "is_gen_extra", False)
        self.dtype_number = self._get_dtype_number()
        self.extra_len_lst = None
        self.length = self._get_case_numbers() if not self.is_gen_extra else self._cal_extra_length()
        self.input_gen = None
        self.method_gens = []
        self.gens = []
        self.create_generates()
        self.tensor_ele_num = 0
        self.shape_process = ShapeProcessor(self)

    def __iter__(self):
        return self

    def __len__(self):
        return self.length

    def __next__(self):
        if self.index >= self.length:
            raise StopIteration
        return self.generate()

    def after_input_config(
            self,
            index: int,
            input_case: Union[InputCaseConfig, List[InputCaseConfig]],
    ) -> Union[InputCaseConfig, List[InputCaseConfig]]:
        logging.debug(f"after_input_config {self.index}")
        return input_case

    def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
        logging.debug(f"after_case_config {self.index}")
        return case_config

    def create_generates(self):
        if self.config.method_inputs:
            for config in self.config.method_inputs:
                instance = self._get_gen_instance(config)
                self.method_gens.append(instance)
        if self.config.tensor_input:
            self.input_gen = self._get_gen_instance(self.config.tensor_input)
        for config in self.config.inputs:
            instance = self._get_gen_instance(config)
            self.gens.append(instance)

    def cal_tensor_ele_num(self, case: InputCaseConfig):
        if not isinstance(case, list) and case.type == "tensor":
            if case.range_values not in ["null", ["null"]]:
                self.tensor_ele_num += cal_tensor_numel(case.shape)
        if isinstance(case, list) and case[0].type == "tensors":
            for tensor in case:
                if tensor.range_values not in ["null", ["null"]]:
                    self.tensor_ele_num += cal_tensor_numel(tensor.shape)

    def process_shape_restrict(
            self,
            all_inputs: List,
            inputs: List,
            method_inputs: List,
            tensor_case: InputCaseConfig,
    ):
        """
        基于用例的所有输入的元组形式，生成每个输入的泛化用例，同时基于最大元素个数的约束，reshape所有输入
        """
        self.tensor_ele_num = 0
        for i, _ in enumerate(self.config.inputs):
            case = self.gens[i].__next__()
            self.cal_tensor_ele_num(case)
            inputs.append(case)
        all_inputs.extend(inputs)

        if self.config.method_inputs:
            for i, _ in enumerate(self.config.method_inputs):
                methods_case = self.method_gens[i].__next__()
                self.cal_tensor_ele_num(methods_case)
                method_inputs.append(methods_case)
        all_inputs.extend(method_inputs)

        if self.config.tensor_input:
            self.cal_tensor_ele_num(tensor_case)
            all_inputs.append(tensor_case)

        self.shape_process.run(all_inputs)

    def generate(self) -> CaseConfig:
        """
        generate case input and attr values,
        you can overwrite it if you need.
        :return: CaseConfig
        """
        all_inputs = []
        inputs = []
        method_inputs = []
        tensor_case = self.input_gen.__next__() if self.input_gen else None
        # process_shape_restrict基于三元组构造泛化用例 由于使用了while 可能会死循环
        # 即使用户重写了after_input_config 仍然可能死循环
        # 想要避免，用户需要重写generate
        # 在生成每个tensor的泛化形式后，tensor之间的约束在after_input_config中体现
        self.process_shape_restrict(all_inputs, inputs, method_inputs, tensor_case)

        case = self._create_case_config()

        index = 0
        if self.config.tensor_input is not None:
            tensor_input = self.after_input_config(index, tensor_case)
            case.tensor_input = tensor_input
            index += 1

        for i, input_case in enumerate(method_inputs):
            method_inputs[i] = self.after_input_config(index, input_case)
            index += 1

        for i, input_case in enumerate(inputs):
            inputs[i] = self.after_input_config(index, input_case)
            index += 1

        case.method_inputs = method_inputs if method_inputs else None
        case.inputs = inputs
        self.index += 1
        return self.after_case_config(case)

    def _get_dtype_number(self):
        return self.config.dtype_numbers

    def _get_case_numbers(self):
        """获取用户指定的用例的数量"""
        # 获取inputs中每个输入的dtype数量，取最大值再乘dtype_numbers
        dtype_len = [len(config.dtypes.values) for config in self.config.inputs]
        if self.config.tensor_input:
            dtype_len.append(len(self.config.tensor_input.dtypes.values))
        if self.config.method_inputs:
            method_dtype_len = [len(config.dtypes.values) for config in self.config.method_inputs]
            dtype_len.extend(method_dtype_len)
        return max(dtype_len) * self.dtype_number

    def _get_gen_instance(self, config):
        # 生成List[tuple]，每个tuple为一个case的元组
        if config.type in PARAMETER_REGISTRY.get_register_keys():
            if self.is_gen_extra and "tensor" in config.type:
                parameter_type_customer = ParameterFactory.create_customer_parameter("extra_" + config.type)
                return parameter_type_customer(config, self.extra_len_lst)
            else:
                parameter_type_customer = ParameterFactory.create_customer_parameter(config.type)
                return parameter_type_customer(config, self.dtype_number, self.length)
        raise ValueError(f"parameter type: {config.type} is not supported, please add custom type")

    def _cal_extra_length(self):
        """
        计算额外用例的数目
        """
        number_process = NumberProcessor(self.config.extra_numbers)
        if self.config.method_inputs:
            for config in self.config.method_inputs:
                if "tensor" in config.type:
                    number_process.append_input(config)
        if self.config.tensor_input is not None:
            if "tensor" in self.config.tensor_input.type:
                number_process.append_input(self.config.tensor_input)
        for config in self.config.inputs:
            if "tensor" in config.type:
                number_process.append_input(config)
        # 计算得到最小数组
        # 格式为[scalar_cases, lower_border_cases, upper_border_cases, empty_cases, infnan_cases]5种tensor的最小数目
        self.extra_len_lst = number_process.cal_length()
        return sum(self.extra_len_lst)

    def _create_case_config(self):
        is_boundary = True if self.is_gen_extra else False
        case = CaseConfig(
            name=self.config.name,
            aclnn_name=self.config.aclnn_name,
            api=self.config.api,
            api_type=self.config.api_type,
            aclnn_api_type=self.config.aclnn_api_type,
            version=self.config.version,
            expected_error_msg=self.config.expected_error_msg,
            backward=self.config.backward,
            standard=self.config.standard,
            outputs=self.config.outputs,
            is_boundary=is_boundary
        )
        return case