import numpy as np
import torch
import torch.nn.functional as F
from modules.basic.utils.registry import METRIC_REGISTRY


def f_score(inputs, target, beta=1, smooth=1e-5, threhold=0.5):
    n, c, h, w = inputs.size()
    nt, ht, wt, ct = target.size()
    if h != ht and w != wt:
        inputs = F.interpolate(
            inputs, size=(ht, wt), mode="bilinear", align_corners=True
        )

    temp_inputs = torch.softmax(
        inputs.transpose(1, 2).transpose(2, 3).contiguous().view(n, -1, c), -1
    )
    temp_target = target.view(n, -1, ct)

    # --------------------------------------------#
    #   计算dice系数
    # --------------------------------------------#
    temp_inputs = torch.gt(temp_inputs, threhold).float()
    tp = torch.sum(temp_target[..., :-1] * temp_inputs, axis=[0, 1])
    fp = torch.sum(temp_inputs, axis=[0, 1]) - tp
    fn = torch.sum(temp_target[..., :-1], axis=[0, 1]) - tp

    score = ((1 + beta**2) * tp + smooth) / (
        (1 + beta**2) * tp + beta**2 * fn + fp + smooth
    )
    score = torch.mean(score)
    return score


# 设标签宽W，长H
def fast_hist(a, b, n):
    # --------------------------------------------------------------------------------#
    #   a是转化成一维数组的标签，形状(H×W,)；b是转化成一维数组的预测结果，形状(H×W,)
    # --------------------------------------------------------------------------------#
    k = (a >= 0) & (a < n)
    # --------------------------------------------------------------------------------#
    #   np.bincount计算了从0到n**2-1这n**2个数中每个数出现的次数，返回值形状(n, n)
    #   返回中，写对角线上的为分类正确的像素点
    # --------------------------------------------------------------------------------#
    return np.bincount(n * a[k].astype(int) + b[k], minlength=n**2).reshape(n, n)


def per_class_iu(hist):
    return np.diag(hist) / np.maximum((hist.sum(1) + hist.sum(0) - np.diag(hist)), 1)


def per_class_PA_Recall(hist):
    return np.diag(hist) / np.maximum(hist.sum(1), 1)


def per_class_Precision(hist):
    return np.diag(hist) / np.maximum(hist.sum(0), 1)


def per_Accuracy(hist):
    return np.sum(np.diag(hist)) / np.maximum(np.sum(hist), 1)


def generate_matrix(gt_image, pre_image, num_classes):
    mask = (gt_image >= 0) & (gt_image < num_classes)
    label = num_classes * gt_image[mask].astype("int") + pre_image[mask]
    count = np.bincount(label, minlength=num_classes**2)
    confusion_matrix = count.reshape(num_classes, num_classes)
    return confusion_matrix


@METRIC_REGISTRY.register()
def calculate_miou(img, img2, img3, img4, num_classes, **kwargs):
    """
    :param img: sr result
    :param img2: sr ground truth
    :param img3: segmentation result
    :param img4: segmentation label
    :param crop_border:
    :param input_order:
    :param kwargs:
    :return:
    """
    confusion_matrix = np.zeros((num_classes,) * 2)
    confusion_matrix += generate_matrix(img4, img3, num_classes)
    MIoU = np.diag(confusion_matrix) / (
        np.sum(confusion_matrix, axis=1)
        + np.sum(confusion_matrix, axis=0)
        - np.diag(confusion_matrix)
    )
    MIoU = np.nanmean(MIoU)
    return MIoU
