"""
Confusion-matrix-based metrics for binary segmentation.

We accumulate TP/FP/FN/TN over an entire epoch (rather than averaging
per-batch ratios) so the reported numbers match the standard global
definitions of IoU, mIoU, Dice, and pixel accuracy.

Reported per call to ``compute()``:
    iou        — foreground IoU = TP / (TP + FP + FN)
    miou       — mean of (foreground IoU, background IoU)
    dice       — 2·TP / (2·TP + FP + FN)
    pixel_acc  — (TP + TN) / total
"""
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, target):
        pred = (torch.sigmoid(logits) > 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)
        miou = 0.5 * (fg_iou + bg_iou)
        dice = (2 * tp) / (2 * tp + fp + fn + eps)
        pixel_acc = (tp + tn) / (tp + tn + fp + fn + eps)

        return {
            "iou": float(fg_iou),
            "miou": float(miou),
            "dice": float(dice),
            "pixel_acc": float(pixel_acc),
        }