from __future__ import annotations

import torch

from src.losses import soft_skeletonize


def binary_metrics(logits: torch.Tensor, target: torch.Tensor, threshold: float = 0.5) -> list[dict[str, float]]:
    pred = torch.sigmoid(logits) >= threshold
    target_b = target >= 0.5
    dims = (1, 2, 3)
    tp = (pred & target_b).float().sum(dim=dims)
    fp = (pred & ~target_b).float().sum(dim=dims)
    fn = (~pred & target_b).float().sum(dim=dims)
    inter = tp
    union = (pred | target_b).float().sum(dim=dims)
    precision = tp / (tp + fp).clamp_min(1.0)
    recall = tp / (tp + fn).clamp_min(1.0)
    dice = (2.0 * inter) / (pred.float().sum(dim=dims) + target_b.float().sum(dim=dims)).clamp_min(1.0)
    iou = inter / union.clamp_min(1.0)
    rows = []
    cldice_scores = cldice_metric(logits, target, threshold)
    for idx in range(logits.size(0)):
        rows.append(
            {
                "iou": float(iou[idx].item()),
                "dice": float(dice[idx].item()),
                "precision": float(precision[idx].item()),
                "recall": float(recall[idx].item()),
                "cldice": float(cldice_scores[idx].item()),
            }
        )
    return rows


def cldice_metric(logits: torch.Tensor, target: torch.Tensor, threshold: float = 0.5) -> torch.Tensor:
    pred = (torch.sigmoid(logits) >= threshold).float()
    target = (target >= 0.5).float()
    pred_skel = soft_skeletonize(pred, 20)
    target_skel = soft_skeletonize(target, 20)
    dims = (1, 2, 3)
    eps = 1.0
    topology_precision = ((pred_skel * target).sum(dim=dims) + eps) / (pred_skel.sum(dim=dims) + eps)
    topology_sensitivity = ((target_skel * pred).sum(dim=dims) + eps) / (target_skel.sum(dim=dims) + eps)
    return (2.0 * topology_precision * topology_sensitivity) / (topology_precision + topology_sensitivity + 1e-6)