# -------------------------------------------------------------------------
#  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, List

import numpy as np
import torch


def compute_error_balance(local_output: torch.Tensor, remote_output: torch.Tensor, small_value=1.0) -> float:
    diff_value = torch.subtract(local_output, remote_output)
    scalar = torch.full_like(remote_output, small_value)
    diff_value_rel = diff_value / torch.max(torch.abs(remote_output), scalar)
    return float(torch.mean(diff_value_rel))


def compute_absolute_error(local_output: torch.Tensor, remote_output: torch.Tensor, small_value=1e-7) -> torch.Tensor:
    return torch.abs(torch.subtract(local_output, remote_output))


def compute_relative_error(local_output: torch.Tensor, remote_output: torch.Tensor, small_value=1e-7) -> torch.Tensor:
    diff_value = torch.abs(torch.subtract(local_output, remote_output))
    diff_value_rel = diff_value / (torch.abs(remote_output) + small_value)
    return diff_value_rel


def compute_root_mean_squared_error(local_output: torch.Tensor, remote_output: torch.Tensor) -> torch.Tensor:
    diff_value = torch.subtract(local_output, remote_output)
    root_mean_squared_error = torch.sqrt(torch.sum(diff_value * diff_value) / diff_value.numel()).item()
    return root_mean_squared_error


def check_invalid_value(value: torch.Tensor):
    has_nan = torch.isnan(value).any()
    has_inf = torch.isinf(value).any()
    return has_nan or has_inf


def compare_torch_data(data, compare_data):
    # 如果类型不同，直接返回 False
    if not isinstance(data, type(compare_data)):
        return False
    # 如果是列表或元组，递归比较每个元素
    if isinstance(data, list) or isinstance(data, tuple):
        if len(data) != len(compare_data):
            return False
        return all(compare_torch_data(d, d1) for d, d1 in zip(data, compare_data))
    # 如果是字典，递归比较每个键值对
    elif isinstance(data, dict):
        if set(data.keys()) != set(compare_data.keys()):
            return False
        return all(compare_torch_data(data[k], compare_data[k]) for k in data)
    # 如果是torch.Tensor，比较张量内容
    elif isinstance(data, torch.Tensor):
        data = data.cpu()
        compare_data = compare_data.cpu()
        if data.dtype == torch.complex32:
            data = data.to(torch.complex64)
            compare_data = compare_data.to(torch.complex64)
        nan_mask_data = torch.isnan(data)
        nan_mask_cmp = torch.isnan(compare_data)
        if not torch.equal(nan_mask_data, nan_mask_cmp):
            return False
        mask = ~nan_mask_data
        return torch.equal(data[mask].cpu(), compare_data[mask].cpu())
    # 如果是布尔值或其他简单类型，直接比较
    else:
        return data == compare_data