# -------------------------------------------------------------------------
#  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 Optional
from pydantic import BaseModel

from atk.common.log import Logger

logging = Logger().get_logger()


class PerfStandardConfig(BaseModel):
    # aver_radio_bm：表示该api的 “所有用例的平均性能比” 需大于该阈值
    # radio_bm：表示该api的 “每个用例的性能比” 均需大于该阈值
    aver_radio_bm: Optional[float] = None
    radio_bm: Optional[float] = None

    def update(self, **kwargs):
        if kwargs.get("aver_radio_bm"):
            self.aver_radio_bm = kwargs.get("aver_radio_bm")
        if kwargs.get("radio_bm"):
            self.radio_bm = kwargs.get("radio_bm")


class AccStandard:
    def __init__(self, dtype, op_type="compute"):
        self.default_standard = DefaultStandards()
        self.standard_dict = self.default_standard.accuracy
        if op_type == "compute":
            dtype_str = str(dtype).split(".")[-1]
            if dtype_str in self.standard_dict:
                self.standard = self.standard_dict[dtype_str]
            else:
                raise ValueError(f"Unsupported dtype: {dtype}")
        else:
            self.standard = [0.0, 0.0, 0.1]

    def get_standard(self):
        if len(self.standard) == 3:
            return self.standard
        else:
            raise ValueError(f"Incompleted accuary standards")


class PerfStandard:
    _instance = None

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(PerfStandard, cls).__new__(cls)
        cls._default_standard = DefaultStandards()
        cls._instance._config = cls._default_standard.performance
        cls._memory_standard = cls._default_standard.memory_standard["pass_ratio"]
        return cls._instance

    @property
    def config(self):
        return self._config

    @property
    def memory_standard(self):
        return self._memory_standard

    def get_standard(self, device: str = "800T_A2", opp_perf: str = "not_key") -> PerfStandardConfig:
        perf_config = PerfStandardConfig()
        device = "800T_A2" if device is None else device
        try:
            if isinstance(opp_perf, str):
                compute_power = self.config["device_compute_map"][device]
                cur_compute_config = self.config[compute_power]
            elif isinstance(opp_perf, list) and len(opp_perf) == 2:
                perf_config.aver_radio_bm = opp_perf[0]
                perf_config.radio_bm = opp_perf[1]
                return perf_config
            else:
                logging.warning(f"no support opp_perf, now will use default config [1.0, 1.0]")
                cur_compute_config = self.config["default"]
        except KeyError as exception:
            logging.warning(f"no perf standard exist for current device: {device}, please add in DefaultStandards, "
                            f"now will use default config [1.0, 1.0]")
            cur_compute_config = self.config["default"]

        perf_config.aver_radio_bm = cur_compute_config[opp_perf][0]
        perf_config.radio_bm = cur_compute_config[opp_perf][1]
        return perf_config


class PassStandard:
    def __init__(self, **kwargs):
        self.default_standard = DefaultStandards()
        self.memory_pass_ratio = self.default_standard.pass_ratio["memory_pass"]
        self.acc_pass_ratio = self.default_standard.pass_ratio["acc_pass"]
        if kwargs.get("acc_pass_ratio"):
            self.acc_pass_ratio = kwargs.get("acc_pass_ratio")

    def get_memory_pass_ratio(self):
        return self.memory_pass_ratio

    def get_acc_pass_ratio(self):
        return self.acc_pass_ratio


class DefaultStandards:
    def __init__(self):
        # 保持原YAML的层级结构
        self.config = {
            'Performance': {
                'default': {
                    'key': [0.95, 0.98],
                    'not_key': [0.95, 0.98],
                    'DSA': [0.95, 0.98]
                },
                'device_compute_map': {
                    '800T_A2': 'default'
                }
            },
            'Accuracy': {
                'float32': [0.0001, 0.0001, 0.1],
                'float64': [0.0001, 0.0001, 0.1],
                'int32': [0.0001, 0.0001, 0.1],
                'int64': [0.0001, 0.0001, 0.1],
                'float16': [0.001, 0.001, 0.1],
                'bfloat16': [0.004, 0.004, 0.1],
                'hfloat32': [0.001, 0.001, 0.1],
                'bool': [0.0, 0.0, 0.1],
                'uint8': [1, 0.01, 0.1],
                'int8': [1, 0.001, 0.1],
                'int16': [1, 0.001, 0.1]
            },
            'Memory_Standard': {
                'pass_ratio': 1.1
            },
            'Pass_Ratio': {
                'memory_pass': 1,
                'acc_pass': 1,
                'mem_check_pass': 1
            }
        }

        # 创建快捷访问属性，简化常用配置的访问
        self.performance = self.config['Performance']
        self.accuracy = self.config['Accuracy']
        self.memory_standard = self.config['Memory_Standard']
        self.pass_ratio = self.config['Pass_Ratio']

    def get(self, *keys, default=None) -> Optional[object]:
        """
        通用获取方法，支持多级键访问
        示例: get('Performance', 'default', 'key')
        """
        current = self.config
        for key in keys:
            if isinstance(current, dict) and key in current:
                current = current[key]
            else:
                return default
        return current