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

from atk.common.log import Logger
from atk.configs.results_config import AccuracyConfig

logging = Logger().get_logger()

NONE_VALUE = None


class SingleBenchmarkCompareStandard:
    def __init__(self):
        self.threshold = {
            torch.float16: 2 ** -10,
            torch.bfloat16: 2 ** -7,
            torch.float32: 2 ** -13,
            torch.float8_e4m3fn: 2 ** -3,
            torch.float8_e5m2: 2 ** -2,
        }
        self.mare_ratio = 10

    def update(self, **kwargs):
        attr_map = {
            "mare_ratio": None,
            "fp16_thd": self.threshold.setdefault,
            "bf16_thd": self.threshold.setdefault,
            "fp32_thd": self.threshold.setdefault,
            "fp8e4m3_thd": self.threshold.setdefault,
            "fp8e5m2_thd": self.threshold.setdefault,
        }
        for key, value in kwargs.items():
            if key in attr_map:
                target = attr_map[key]
                if target is None:
                    setattr(self, key, value)
                else:
                    target(torch.__dict__[key.removesuffix("_thd").replace("fp8e4m3", "float8_e4m3fn") \
                           .replace("fp8e5m2", "float8_e5m2").replace("fp16", "float16") \
                           .replace("bf16", "bfloat16").replace("fp32", "float32")], value)

    def get_threshold(self, dtype):
        if dtype == torch.float64:
            logging.warning("The output data of fp64 uses the same standard as fp32.")
            return self.threshold.get(torch.float32)
        if dtype in self.threshold.keys():
            return self.threshold.get(dtype)
        logging.error("Single benchmark compare only supports floating point in fp16, bf16, fp32, fp8e4m3, fp8e5m2.")
        return NONE_VALUE


class SingleBenchSummary:
    def __init__(self, acc_result: AccuracyConfig, threshold, mare_ratio):
        self.check_result = None
        self.mean_rel_err = acc_result.mean_rel_err
        self.max_rel_err = acc_result.max_rel_err
        self.threshold = threshold
        self.mare_ratio = mare_ratio
        self._get_check_result()

    def get_result_msg(self):
        result_str = ""
        if self.check_result:
            return result_str
        if self.mean_rel_err >= self.threshold:
            result_str += f"平均相对误差MERE为{self.mean_rel_err}, 超过阈值{self.threshold}.\n"
        if self.max_rel_err >= self.threshold * self.mare_ratio:
            result_str += f"最大相对误差MARE为{self.max_rel_err}, 超过阈值{self.threshold * self.mare_ratio}.\n"
        return result_str

    def _get_check_result(self):
        if self.mean_rel_err >= self.threshold or self.max_rel_err >= self.threshold * self.mare_ratio:
            self.check_result = False
        else:
            self.check_result = True