# -------------------------------------------------------------------------
#  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.
# -------------------------------------------------------------------------

import sys
from decimal import Decimal
from functools import reduce
from typing import List, Union

from atk.case_generator.generator.parameter_types import ParameterFactory
from atk.case_generator.utils.case_utils import ExtraInputLength
from atk.common.log import Logger
from atk.configs.design_config import InputDesignConfig

logging = Logger().get_logger()
EXTRA_ENUM = ExtraInputLength.get_enum()


class ShapeProcessor:
    def __init__(self, case_gen, max_reshape_cnt: int = 100):
        self.case_gen = case_gen
        self.config = case_gen.config
        self.max_reshape_cnt = max_reshape_cnt
        self.shape_intervals = []
        self.shape_proportions = []
        self.restrict_nums = []
        self.load_intervals_and_proportions()

    @staticmethod
    def decimal_func(opt1: str, opt2: str, cal: str) -> float:
        decimal_opt1 = Decimal(str(opt1))
        decimal_opt2 = Decimal(str(opt2))

        if cal == "+":
            result = decimal_opt1 + decimal_opt2
        elif cal == "-":
            result = decimal_opt1 - decimal_opt2
        else:
            raise ValueError("Invalid operation. Use '+' or '-' for cal argument.")

        return float(result)

    def run(self, inputs: List) -> None:
        """
        Tensor shape元素分布比例配置的执行入口
        调用此方法实现元素分布比例可配置化功能
        @param max_reshape_cnt:
        @param inputs: 当前case的inputs列表
        @return:
        """
        reshape_cnt = 0

        while reshape_cnt < self.max_reshape_cnt:
            is_put = self.execute_tensor_restrict(is_put=False)
            if is_put:
                break
            self.regenerate_for_tensor_shape(inputs)
            reshape_cnt += 1

    def execute_tensor_restrict(self, is_put: bool) -> bool:
        """
        执行tensor元素限制策略
        若当前case生成的元素个数在所处区间内已达到restrict_nums最大容量上限，则不放行case生成通过
        @param is_put: 标志位，表示当前case是否满足放入对应区间的要求
        @return:
        """
        for idx, [start, end] in enumerate(self.shape_intervals):
            if start <= self.case_gen.tensor_ele_num < end:
                is_put = self.execute_put_and_change(is_put, idx)
        return is_put

    def regenerate_for_tensor_shape(self, inputs: List) -> None:
        """
        重新随机生成当前case的shape元素个数
        @param inputs: 当前case的inputs列表
        @return:
        """
        self.case_gen.tensor_ele_num = 0
        offset = 0
        for index, gen in enumerate(self.case_gen.gens):
            self.reshape_and_cal_tensor_ele_num(inputs, gen, offset, index)

        offset += len(self.case_gen.gens)
        for index, gen in enumerate(self.case_gen.method_gens):
            self.reshape_and_cal_tensor_ele_num(inputs, gen, offset, index)

        offset += len(self.case_gen.method_gens)
        if self.case_gen.input_gen:
            self.reshape_and_cal_tensor_ele_num(
                inputs, self.case_gen.input_gen, offset, 0
            )

    def reshape_and_cal_tensor_ele_num(self, inputs, gen, offset, index):
        if isinstance(gen, ParameterFactory.create_customer_parameter("tensor")):
            dim = (
                0
                if inputs[offset + index].shape == [0]
                else len(inputs[offset + index].shape)
            )
            inputs[offset + index].shape = gen.get_shapes_by_dim(
                dim, inputs[offset + index].dtype
            )
            if inputs[offset + index].range_values in ["null", ["null"]]:
                return
            self.case_gen.tensor_ele_num += reduce(
                lambda x, y: x * y, inputs[offset + index].shape, 1
            )
        if isinstance(gen, ParameterFactory.create_customer_parameter("tensors")):
            for gen_case in inputs[offset + index]:
                dim = 0 if gen_case.shape == [] else len(gen_case.shape)
                gen_case.shape = gen.gen.get_shapes_by_dim(dim, gen_case.dtype)
                if gen_case.range_values in ["null", ["null"]]:
                    return
                self.case_gen.tensor_ele_num += reduce(
                    lambda x, y: x * y, gen_case.shape, 1
                )

    def process_large_shape_config(
            self,
            shape_config: List[List[Union[int, float]]],
            ori_intervals: List,
            ori_proportions: List,
            intervals: List,
            proportions: List,
    ) -> None:
        if not shape_config:
            return
        shape_config.sort(key=lambda x: x[0], reverse=True)
        start, end = ori_intervals[0]
        remain_per = ori_proportions[0]
        calculated_per = 0
        large_intervals = []
        large_proportions = []
        for large_rule in shape_config:
            if large_rule[0] >= sys.maxsize:
                raise ValueError(
                    "Invalid scope definition in shape restriction configuration"
                )
            large_intervals.insert(0, [large_rule[0], end])
            large_proportions.insert(
                0, self.decimal_func(large_rule[1], calculated_per, "-")
            )
            remain_per = self.decimal_func(ori_proportions[0], large_rule[1], "-")
            calculated_per = large_rule[1]
            end = large_rule[0]
        large_intervals.insert(0, [start, end])
        large_proportions.insert(0, remain_per)
        intervals.extend(large_intervals)
        proportions.extend(large_proportions)

    def load_intervals_and_proportions(self):
        shape_config = self.config.shape_distributions
        ori_intervals = [[0, sys.maxsize]]
        ori_proportions = [1.0]
        intervals = []
        proportions = []
        try:
            self.process_large_shape_config(
                shape_config, ori_intervals, ori_proportions, intervals, proportions
            )
        except ValueError as exception:
            logging.error(exception)
        self.shape_intervals = intervals
        self.shape_proportions = proportions
        self.restrict_nums = [
            [0, int(proportion * self.case_gen.length + 0.5)]
            for proportion in self.shape_proportions
        ]
        diff = self.case_gen.length - sum(
            [restrict_num[1] for restrict_num in self.restrict_nums]
        )
        if diff != 0:
            for i in range(diff):
                self.restrict_nums[i % len(self.restrict_nums)][1] += 1

    def execute_put_and_change(self, is_put: bool, idx: int) -> bool:
        if self.restrict_nums[idx][0] + 1 <= self.restrict_nums[idx][1]:
            self.restrict_nums[idx][0] += 1
            is_put = True
            self.change_tensor_max_ele(lambda x: x.config.shapes.max_length)
        else:
            need_to_change = True
            for index in range(idx, len(self.restrict_nums)):
                if self.restrict_nums[index][0] < self.restrict_nums[index][1]:
                    need_to_change = False
            if need_to_change:
                self.change_tensor_max_ele(lambda x: int(x.max_number_ele / 2))
        return is_put

    def change_tensor_max_ele(self, func):
        for gen in self.case_gen.gens:
            if isinstance(gen, ParameterFactory.create_customer_parameter("tensor")):
                gen.update(max_number_ele=func(gen))
            if isinstance(gen, ParameterFactory.create_customer_parameter("tensors")):
                gen.gen.update(max_number_ele=func(gen.gen))
        for gen in self.case_gen.method_gens:
            if isinstance(gen, ParameterFactory.create_customer_parameter("tensor")):
                gen.update(max_number_ele=func(gen))
            if isinstance(gen, ParameterFactory.create_customer_parameter("tensors")):
                gen.gen.update(max_number_ele=func(gen.gen))
        if self.case_gen.input_gen:
            if isinstance(
                    self.case_gen.input_gen,
                    ParameterFactory.create_customer_parameter("tensor"),
            ):
                self.case_gen.input_gen.update(
                    max_number_ele=func(self.case_gen.input_gen)
                )
            if isinstance(
                    self.case_gen.input_gen,
                    ParameterFactory.create_customer_parameter("tensors"),
            ):
                self.case_gen.input_gen.gen.update(
                    max_number_ele=func(self.case_gen.input_gen.gen)
                )

    def set_max_reshape_cnt(self, new_cnt):
        self.max_reshape_cnt = new_cnt


class NumberProcessor:
    """
    用于让yaml中不同tensor输入的不同种类的特殊tensor数量对齐
    特殊tensor包括5种: scalar_cases, lower_border_cases, upper_border_cases, empty_cases, infnan_cases
    由于用例生成的原则是同一个用例里, 特殊tensor应当为同一类特殊,
    因此需保证不同输入的同类tensor数量需对齐
    该类先收集各个输入的各类特殊tensor的数目, 再取不同输入中, 同类特殊数目的最大数目作为该类特殊tensor的数目
    如果遇到该输入的has_empty或has_infnan等参数为False, 则不参与计算最大值
    """

    def __init__(self, extra_numbers: Union[str, int] = 0):
        if extra_numbers == 'all' or (isinstance(extra_numbers, int) and extra_numbers >= 0):
            self.numbers = extra_numbers
        else:
            raise ValueError("extra_numbers should be a non-negative int or string 'all'")
        self.len_lsts: List[ExtraInputLength] = []
        self.flag_lst = {extra_enum: [] for extra_enum in EXTRA_ENUM}
        self.final_len_lst: List[int] = []

    def append_input(self, config: InputDesignConfig) -> None:
        #  len_lst: [scalar_cases, lower_border_cases, upper_border_cases, empty_cases, infnan_cases]5种tensor的数目列表
        create_parameter_type = ParameterFactory.create_customer_parameter(
            "extra_" + config.type
        )
        len_lst = create_parameter_type.cal_length(config)
        self.len_lsts.append(len_lst)
        for extra_enum in EXTRA_ENUM:
            self.flag_lst[extra_enum].append(getattr(config.boundary, 'has_' + extra_enum))

    def cal_length(self):
        if not self.len_lsts:
            return []
        bool_lsts = list(self.flag_lst.values())

        for i in range(5):
            # 参与计算的列表，5类特殊tensor均可能生成
            bool_lst = bool_lsts[i]
            participate_lst = [
                self.len_lsts[j][i]
                for j, bool_value in enumerate(bool_lst)
                if bool_value
            ]
            self.final_len_lst.append(max(participate_lst) if participate_lst else 0)

        if not isinstance(self.final_len_lst, list) or len(self.final_len_lst) != 5:
            raise ValueError(f"Length of final_len_lst should be 5, please check!")

        if isinstance(self.numbers, int):
            self._adjust_len()
        return ExtraInputLength(*self.final_len_lst)

    def _adjust_len(self):
        # 更新final_len_lst，如果用户未指定dtype_extra_numbers，则按照遍历的数目来，否则按照用户指定的数目来生成
        diff = self.numbers - sum(self.final_len_lst)
        if diff > 0:
            # 如果差值大于0，按照比例增加数值
            ratio = [x / sum(self.final_len_lst) for x in self.final_len_lst]
            incre = [int(diff * r) for r in ratio]
            self.final_len_lst = [x + inc for x, inc in zip(self.final_len_lst, incre)]
            # 将剩余的差值分配到前面的元素
            for i in range(self.numbers - sum(self.final_len_lst)):
                self.final_len_lst[i] += 1
        elif diff < 0:
            # 如果差值小于0，按照比例减少数值，确保不出现负值
            ratio = [x / sum(self.final_len_lst) for x in self.final_len_lst]
            decre = [int(abs(diff) * r) for r in ratio]
            # 确保减少的值不超过当前元素的值
            for i, length in enumerate(self.final_len_lst):
                if decre[i] > length:
                    decre[i] = length
            self.final_len_lst = [x - dec for x, dec in zip(self.final_len_lst, decre)]
            # 将剩余的差值从前面的元素中抽走
            for i in range(sum(self.final_len_lst) - self.numbers):
                if self.final_len_lst[i] > 0:
                    self.final_len_lst[i] -= 1