import matplotlib.pyplot as plt
import torch
import torch.nn as nn


def iou(outputs: torch.Tensor, labels: torch.Tensor):
    SMOOTH = 1e-6
    outputs = (outputs.squeeze(1) > 0.5).int()
    labels = (labels > 0.5).int()
    intersection = torch.sum((outputs & labels).float())
    union = torch.sum((outputs | labels).float())

    iou = (intersection + SMOOTH) / (union + SMOOTH)
    return iou


def sis(new_mask, old_mask, new_seg, old_seg):
    def difference(mask1, mask2):
        return torch.round(mask1) * (1 - torch.round(mask2)) + torch.round(mask2) * (
                1 - torch.round(mask1))

    epsilon = 1e-5
    sis = torch.sum(
        difference(
            difference(new_mask, old_mask),
            difference(new_seg, old_seg))
    ) / torch.sum(torch.clamp(new_mask + old_mask + new_seg + old_seg, 0, 1) + epsilon)  # normalizing factor
    return sis


def precision(output, labels, threshold):
    t = (output > threshold).float()
    tp = torch.sum(t * labels)
    fp = torch.sum(t * (1 - labels))
    return tp / (tp + fp + 1e-5)


def recall(output, labels, threshold):
    t = (output > threshold).float()
    tp = torch.sum(t * labels)
    fn = torch.sum((1 - t) * labels)
    return tp / (tp + fn + 1e-5)


def tp_rate(output, labels, threshold):
    t = (output > threshold).float()
    tp = torch.sum(t * labels)
    fn = torch.sum((1 - t) * labels)
    return tp / (tp + fn + 1e-5)


def fp_rate(output, labels, threshold):
    t = (output > threshold).float()
    fp = torch.sum(t * (1 - labels))
    tn = torch.sum((1 - t) * (1 - labels))
    return fp / (fp + tn + 1e-5)


if __name__ == '__main__':
    test = torch.zeros(10, 10)
    test[:3, :3] = 1
    test2 = torch.zeros(10, 10)
    test2[:3, :3] = 1
    print(iou(test, test2))