"""Confusion-matrix-based metrics. Same as the other experiments."""
import torch


class SegMetrics:
    def __init__(self, threshold=0.5):
        self.threshold = threshold
        self.reset()

    def reset(self):
        self.tp = 0.0
        self.fp = 0.0
        self.fn = 0.0
        self.tn = 0.0

    @torch.no_grad()
    def update(self, logits_or_probs, target, output_is_prob=False):
        pred = logits_or_probs if output_is_prob else torch.sigmoid(logits_or_probs)
        pred = (pred > self.threshold).float()
        target = target.float()
        self.tp += (pred * target).sum().item()
        self.fp += (pred * (1 - target)).sum().item()
        self.fn += ((1 - pred) * target).sum().item()
        self.tn += ((1 - pred) * (1 - target)).sum().item()

    def compute(self):
        eps = 1e-7
        tp, fp, fn, tn = self.tp, self.fp, self.fn, self.tn
        fg_iou = tp / (tp + fp + fn + eps)
        bg_iou = tn / (tn + fp + fn + eps)
        return {
            "iou": float(fg_iou),
            "miou": float(0.5 * (fg_iou + bg_iou)),
            "dice": float((2 * tp) / (2 * tp + fp + fn + eps)),
            "pixel_acc": float((tp + tn) / (tp + tn + fp + fn + eps)),
        }