# -------------------------------------------------------------------------
#  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 logging
import random
import math
from enum import Enum
from typing import Any, Optional, List
from typing import Union, Dict
from typing_extensions import Literal

import yaml
from pydantic import BaseModel, ConfigDict, model_validator, PrivateAttr

DEFAULT_TYPES = ["tensor", "tensors", "tensor_tuple",
                 "scalar", "scalars", "scalar_tuple",
                 "attr", "attrs", "attr_tuple"]

DEFAULT_DTYPES = [
    "fp32",
    "bf16",
    "fp16",
    "fp64",
    "hf32",
    "int8",
    "uint8",
    "int16",
    'uint16',
    "int32",
    "uint32",
    "int64",
    "uint64",
    "bool",
]
ATTR_DTYPES = ["float", "int", "string", "attr_bool"]
NEW_DTYPES = ["complex128", "complex64", "complex32"]

DEFAULT_DIM_NUMBERS = [1, 2, 3, 4, 5, 6, 7, 8]
DEFAULT_DIM_WEIGHTS = [1 / 8] * 8

DEFAULT_DIM_VALUES = [
    1,
    7,
    8,
    9,
    15,
    16,
    17,
    19,
    20,
    21,
    [255, 257],
    131073,
    2147483648,
    [1, 1024],
    [1025, 10240],
    [10241, 102400],
    [102401, 1024000],
    [1024001, 10240000],
    [10240001, 102400000],
    [102400001, 1024000000],
    [1024000001, 2147483648],
]
DEFAULT_DIM_VALUES_WEIGHTS = [
    0.05,
    0.03,
    0.04,
    0.03,
    0.03,
    0.04,
    0.03,
    0.03,
    0.04,
    0.03,
    0.05,
    0.05,
    0.001,
    0.001,
    0.001,
    0.001,
    0.001,
    0.001,
    0.001,
    0.001,
    0.001,
]

DEFAULT_RANGES_VALUES = [[-5, 5]]

DEFAULT_TENSOR_RANGES_INVALID_VALUES = [["-inf"], ["inf"], ["nan"], ["null"]]
DEFAULT_SCALAR_RANGES_INVALID_VALUES = [["-inf"], ["inf"], ["nan"]]
DEFAULT_ATTR_RANGES_INVALID_VALUES = [["-inf"], ["inf"]]

VALID_WEIGHTS = 0.95
MAX_NUMBER_OF_ELEMENTS = 2 ** 32
DEFAULT_DTYPE_NUMBERS = 50


class RandomTypes(Enum):
    DEFAULT = "default"  # 均匀分布
    ND = "nd"  # 正态分布


class RandomTypesConfig(BaseModel):
    model_config = ConfigDict(extra='forbid', ignored_types=(type(lambda x: x),))
    name: Optional[RandomTypes] = RandomTypes.DEFAULT
    mean: Optional[List[float]] = [-100, 100]
    std: Optional[List[float]] = [1, 25]

    @model_validator(mode="after")
    def check(self):
        if self.name == RandomTypes.ND:
            if not self.mean:
                raise ValueError("not set nd mean, please check!")
            if not self.std:
                raise ValueError(" not set nd std, please check!")
        return self

    def model_dump(self):
        data = super(RandomTypesConfig, self).model_dump()
        if data.get("name"):
            data["name"] = data["name"].value
        return data


class RandomConfig(BaseModel):
    model_config = ConfigDict(extra='forbid', ignored_types=(type(lambda x: x),))
    type: Optional[Literal["choices"]] = "choices"
    values: List = None
    weights: Optional[List[float]] = None
    invalid_values: List[Union[str, int]] = []
    random_types: Optional[List[RandomTypesConfig]] = None
    custom_dist_ratio: Optional[Union[List[List[float]], str]] = None  # 表示自定义分布比例

    @model_validator(mode="after")
    def weight_len_must_same_values(self):
        if self.weights and len(self.weights) != len(self.values):
            raise ValueError("length of weights and values is not same.")
        return self

    def get_values(self, k=1, clc_interval=True):
        fun = getattr(self, self.type)
        if k == 1:
            return fun(k, clc_interval)[0]
        else:
            return fun(k, clc_interval)

    def choices(self, k, clc_interval):
        ret = []
        ret_nd = []
        if self.random_types:
            for random_type in self.random_types:
                if random_type.name != RandomTypes.DEFAULT:
                    ret_nd.append(random_type.model_dump())
        for _ in range(k):
            index = random.choices(range(len(self.values)), weights=self.weights, k=k)[0]
            value = self.values[index]
            double_int = (
                    clc_interval
                    and isinstance(value, list)
                    and len(value) == 2
                    and isinstance(value[0], int)
                    and isinstance(value[1], int)
            )
            if double_int:
                ret.append(random.randint(value[0], value[1]))
            else:
                ret.append(value)
            if ret_nd:
                return ret_nd if len(self.random_types) == 1 else random.choice([ret, ret_nd])
        return ret

    def get_actual_values(self):
        # 获得真实values
        ret = []
        if not self.random_types:
            return list(self.values)
        for value in self.values:
            for random_type in self.random_types:
                if random_type.name != RandomTypes.DEFAULT:
                    ret.append(random_type.model_dump())
                else:
                    ret.append(value)
        return ret

    def get_actual_values_with_custom_dist(self, custom_dist_ratio, len_cases: int):
        """
        自定义分布比例
        例如custom_dist_ratio为[[0.1, 0.3, 0.1], [0.4, 0.1]]时，其中第一个列表表示均匀分布相关比例，第二个列表表示正态分布相关比例
        其中，[0.1, 0.3, 0.1]表示第一个均匀分布范围的比例为0.1，第二个均匀分布范围的比例为0.3，最后一个数固定表示离群分布比例，此时为0.1
        [0.4, 0.1]表示正态分布的比例为0.4，离群分布比例为0.1
        :param custom_dist_ratio:
        :param len_cases:
        :return:
        """
        if isinstance(custom_dist_ratio, str) and custom_dist_ratio == "default":
            custom_dist_ratio = DefaultCustomDistConfig().custom_dist_ratio
            self.values = DefaultCustomDistConfig().default_range_values
            logging.info("use default custom_dist_ratio, "
                         f"custom_dist_ratio is {custom_dist_ratio}, self.values is {self.values}")

        ret = []
        # 校验参数合法性
        if not self.random_types or len(custom_dist_ratio) != 2:
            logging.error(f"has no random_types: {self.random_types} or "
                          f"len of custom_dist_ratio is not 2: {custom_dist_ratio}, return default values!!!")
            return list(self.values)

        total_ratio = sum(sum(sublist) for sublist in custom_dist_ratio)
        if abs(total_ratio - 1.0) > 1e-9:
            logging.error(f"custom_dist_ratio is {custom_dist_ratio}, sum is not 1, return default values!!!")
            return list(self.values)

        if len(self.values) != len(custom_dist_ratio[0]) - 1:
            logging.error(f"values length is {len(self.values)}, "
                          f"values_ratio in custom_dist_ratio length is {len(custom_dist_ratio[0]) - 1}")
            return list(self.values)
        # 结束校验

        nd_ratio = custom_dist_ratio[1]  # custom_dist_ratio中定义的是均匀分布（位置0）、正态分布（位置1）范围的比例
        nd_number = round(sum(nd_ratio) * len_cases)  # 正态分布case数量
        if nd_number < 0 or nd_number > len_cases:
            logging.error(f"nd_dist_number illegal, value is {nd_number}, "
                          f"len_cases is {len_cases}, return default values!!!")
            return list(self.values)
        for random_type in self.random_types:
            if random_type.name != RandomTypes.DEFAULT:
                ret.extend([random_type.model_dump()] * nd_number)

        default_number = len_cases - nd_number  # 均匀分布case数量
        default_ratio_total = sum(custom_dist_ratio[0])
        if default_number <= 0 or default_ratio_total <= 0:
            return ret

        for i, value in enumerate(self.values):
            value_ratio = custom_dist_ratio[0][i] / default_ratio_total
            if i == len(self.values) - 1:  # 最后一个value中包含离群值比例，后续再设置离群点
                value_ratio = (custom_dist_ratio[0][i] + custom_dist_ratio[0][-1]) / default_ratio_total
            value_number = math.ceil(value_ratio * default_number)  # 向上取整
            ret.extend([value] * value_number)

        return ret


class DefaultCustomDistConfig(BaseModel):
    default_range_values: List[Union[List[float], float]] = [[-5, 5]]
    custom_dist_ratio: Optional[List[List[float]]] = [[0.5, 0.0], [0.5, 0.0]]
    default_outlier_values: List[float] = [0.001, 1000]  # 表示默认离群值范围
    default_outlier_values_range_index: int = -1  # 表示默认离群值范围在default_range_values中的索引


class InputsShape(BaseModel):
    model_config = ConfigDict(extra='forbid', ignored_types=(type(lambda x: x),))
    dim_numbers: RandomConfig = RandomConfig(
        values=DEFAULT_DIM_NUMBERS, weights=DEFAULT_DIM_WEIGHTS
    )
    dim_values: Optional[Union[RandomConfig, List[RandomConfig]]] = RandomConfig(
        values=DEFAULT_DIM_VALUES, weights=DEFAULT_DIM_VALUES_WEIGHTS
    )
    max_length: int = MAX_NUMBER_OF_ELEMENTS

    @model_validator(mode="after")
    def check_params(self):
        if isinstance(self.dim_values, list):
            dim_values_len = len(self.dim_values)
            max_dim_numbers = max(self.dim_numbers.values)
            if dim_values_len != max_dim_numbers:
                raise ValueError(f"dim_values_length {dim_values_len} is not equal to "
                                 f"max(dim_numbers) {max_dim_numbers}, please check")
        return self


class InputSize(BaseModel):
    model_config = ConfigDict(extra='forbid', ignored_types=(type(lambda x: x),))
    values: List[Union[List[float], float]]
    ratio: float

    @model_validator(mode='after')
    def validate_values_format(self):
        for element in self.values:
            if isinstance(element, list):
                if len(element) != 2:
                    raise ValueError(
                        "Unsupported data type. "
                        f"Use List[2] = [float, float] for range of value, instead of {element}."
                    )
                elif element[0] > element[1]:
                    raise ValueError(
                        f"In range of value {element}, "
                        "the upper bound should be greater than the lower bound."
                    )
        return self


class RangeConfig(BaseModel):
    model_config = ConfigDict(extra='forbid', ignored_types=(type(lambda x: x),))
    valid: RandomConfig = RandomConfig(values=DEFAULT_RANGES_VALUES)
    invalid: RandomConfig = RandomConfig()
    valid_weights: float = VALID_WEIGHTS

    def update(self, **kwargs):
        for key, value in kwargs.items():
            setattr(self, key, value)


class BoundaryConfig(BaseModel):
    model_config = ConfigDict(extra='forbid', ignored_types=(type(lambda x: x),))
    has_empty: Optional[bool] = True
    has_infnan: Optional[bool] = True
    has_scalar: Optional[bool] = True
    has_upper_border: Optional[bool] = True
    has_lower_border: Optional[bool] = True

    # 用于记录哪些字段是用户显式设置的
    _user_set_fields: set = PrivateAttr(default_factory=set)

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._user_set_fields = set(kwargs.keys())

    @property
    def user_set_fields(self):
        return self._user_set_fields


class InputDesignConfig(BaseModel):
    model_config = ConfigDict(extra='forbid', ignored_types=(type(lambda x: x),))
    name: Optional[str] = None
    aclnn_name: Optional[str] = None
    type: Optional[str]
    required: bool = True
    backward: Optional[bool] = True
    align_32B: Optional[bool] = None
    outlier_values: Optional[List[float]] = None
    dtypes: Optional[RandomConfig] = RandomConfig(values=DEFAULT_DTYPES)
    ranges: Optional[RangeConfig] = RangeConfig()
    shapes: Optional[InputsShape] = InputsShape()
    tuple_numbers: Optional[RandomConfig] = RandomConfig(values=[1])
    boundary: Optional[BoundaryConfig] = BoundaryConfig()
    expected_error_msg: Optional[str] = None

    @model_validator(mode="after")
    def set_default_ranges_by_type(self):
        tensor_like_types = ["tensor", "tensors", "tensor_tuple", "scalar", "scalars", "scalar_tuple"]
        if self.type in tensor_like_types and "valid" not in self.ranges.model_fields_set:
            self.ranges.valid.values = DefaultCustomDistConfig().default_range_values
            self.ranges.valid.random_types = [RandomTypesConfig(name=RandomTypes.ND)]
            self.ranges.valid.custom_dist_ratio = "default"
        if not self.ranges.invalid.values:
            self.ranges.invalid.values = self.ranges.valid.values
        return self

    @model_validator(mode="after")
    def check_extra_params(self):
        """非tensor类的inputs的边界参数不允许设置, 且默认值为None"""
        if self.type not in ["tensor", "tensors", "tensor_tuple"] and self.boundary:
            # 检查 boundary 中是否有用户显式设置的值
            for field_name in self.boundary.user_set_fields:
                raise ValueError(f"type: '{self.type}' does not support parameters: {field_name}")
        return self

    @model_validator(mode="after")
    def check_type_params(self):
        """检查 type 参数是否合法"""
        if self.type not in DEFAULT_TYPES:
            raise ValueError(f"type: '{self.type}' is not a valid type")
        return self

    @model_validator(mode="after")
    def check_tuple_numbers_param(self):
        """检查 tuple_numbers 参数是否合法"""
        if self.tuple_numbers.values != [1]:
            if self.type not in ["tensors", "tensor_tuple", "scalars", "scalar_tuple", "attrs", "attr_tuple"]:
                raise ValueError(f"type: '{self.type}' does not support tuple_numbers, please check.")
        return self

    @model_validator(mode="after")
    def check_outlier_values_param(self):
        """检查 outlier_values 参数是否合法"""
        if self.outlier_values and len(self.outlier_values) != 2:
            raise ValueError(f"outlier_values: '{self.typoutlier_valuese}' should have two values, please check.")
        return self


CATCCOS_SHAPE_COLUMNS = ("M", "K", "N", "Transpose A", "Transpose B")


class CatccosShapeInput(BaseModel):
    model_config = ConfigDict(extra='forbid')
    name: str
    values: List[Union[int, str]]

    @model_validator(mode="after")
    def check_name(self):
        if self.name not in CATCCOS_SHAPE_COLUMNS:
            raise ValueError(
                f"shape_inputs.name must be one of {CATCCOS_SHAPE_COLUMNS}, got {self.name!r}"
            )
        if not self.values:
            raise ValueError(f"shape_inputs {self.name!r} values must not be empty")
        return self


class StandardConfig(BaseModel):
    model_config = ConfigDict(extra='forbid', ignored_types=(type(lambda x: x),))
    acc: Optional[Union[str, dict]] = "default"
    perf: Optional[Union[str, List[float], Dict[str, List[float]]]] = "not_key"
    mem: Optional[float] = 1.1

    def is_acc_benchmark(self):
        if isinstance(self.acc, str):
            name = self.acc
        else:
            name = list(self.acc.keys())[0]

        return "benchmark" in name

    def is_bm_benchmark(self):
        if isinstance(self.acc, str):
            name = self.acc
        else:
            name = list(self.acc.keys())[0]

        return "bm" in name


class DesignConfig(BaseModel):
    model_config = ConfigDict(extra='forbid', ignored_types=(type(lambda x: x),))

    def __init__(self, path):
        with open(path, "r") as fout:
            config = yaml.safe_load(fout)
        super().__init__(**config)
        self.adapt_shape_distribution()

    name: Optional[str] = None
    aclnn_name: Optional[str] = None
    triton_name: Optional[str] = None
    kernel_name: Optional[str] = None
    version: Optional[str] = None
    expected_error_msg: Optional[str] = None
    api: str = "pytorch"
    api_type: str = "function"
    aclnn_api_type: str = "aclnn_function"
    triton_api_type: str = "triton_function"
    fusion_api_type: str = "fusion_function"
    dist_api_type: str = "dist_function"
    kernel_api_type: str = "kernel_function"
    generate: str = "default"
    standard: Optional[StandardConfig] = StandardConfig()
    backward: bool = False
    outputs: Optional[Union[str, int]] = None
    inputs: Optional[List[InputDesignConfig]] = []
    tensor_input: InputDesignConfig = None
    method_inputs: List[InputDesignConfig] = None
    dt_config: Optional[Dict[str, Any]] = None
    shape_inputs: Optional[List[CatccosShapeInput]] = None
    size_distributions: List[InputSize] = None
    shape_distributions: Optional[List[List[Union[int, float]]]] = [
        [10000000, 0.3],
        [100000000, 0.1],
        [1000000000, 0.01],
    ]
    dtype_numbers: Optional[int] = DEFAULT_DTYPE_NUMBERS
    extra_numbers: Optional[Union[str, int]] = "all"
    compute_times: Optional[int] = None
    _is_gen_extra: Optional[bool] = False

    @property
    def is_gen_extra(self):
        return self._is_gen_extra

    @model_validator(mode="after")
    def check_catccos_fields(self):
        if self.generate != "catccos":
            return self
        if not self.dt_config:
            raise ValueError("generate=catccos requires dt_config")
        if not self.shape_inputs:
            raise ValueError("generate=catccos requires shape_inputs")
        required_dt = ("kernel_name", "data_type", "device_list", "catccos_root")
        missing = [k for k in required_dt if k not in self.dt_config or self.dt_config[k] in (None, "")]
        if missing:
            raise ValueError(f"dt_config missing required fields: {missing}")
        if self.extra_numbers == "all" or (isinstance(self.extra_numbers, int) and self.extra_numbers > 0):
            object.__setattr__(self, "extra_numbers", 0)
        return self

    @model_validator(mode="after")
    def check_dtype_numbers_and_extra_numbers_is_valid(self):
        if self.dtype_numbers < 0:
            raise ValueError("dtype_numbers should be non-negative but is set to negative!")
        if isinstance(self.extra_numbers, int) and self.extra_numbers < 0:
            raise ValueError("extra_numbers should be a non-negative int or string 'all' but is set to negative!")
        if isinstance(self.extra_numbers, str) and self.extra_numbers != 'all':
            raise ValueError("extra_numbers should be a non-negative int or string 'all' but is set to other string!")
        return self

    @is_gen_extra.setter
    def is_gen_extra(self, bool_value):
        self._is_gen_extra = bool_value

    def adapt_shape_distribution(self):
        # catccos 用 shape_inputs，不参与 pytorch 的 shape_distributions 缩放
        if self.generate == "catccos":
            return

        all_shapes = []
        if self.inputs:
            all_shapes.extend(
                [
                    cur_input.shapes
                    for cur_input in self.inputs
                    if isinstance(cur_input, InputDesignConfig)
                ]
            )
        if isinstance(self.tensor_input, InputDesignConfig):
            all_shapes.append(self.tensor_input.shapes)
        if self.method_inputs:
            all_shapes.extend(
                [
                    cur_input.shapes
                    for cur_input in self.inputs
                    if isinstance(cur_input, InputDesignConfig)
                ]
            )

        if not all_shapes:
            return

        total_length = sum(cur_shape.max_length for cur_shape in all_shapes)
        ratio = total_length / (MAX_NUMBER_OF_ELEMENTS * len(all_shapes))
        for shape_distribution in self.shape_distributions:
            shape_distribution[0] = round(shape_distribution[0] * ratio)