from __future__ import annotations
from typing import Dict, List
import numpy as np
import torch
from scipy import ndimage


def dice_per_class(pred: np.ndarray, target: np.ndarray, num_classes: int, include_bg: bool = False) -> Dict[str, float]:
    out = {}
    start = 0 if include_bg else 1
    vals = []
    for c in range(start, num_classes):
        p = pred == c
        t = target == c
        denom = p.sum() + t.sum()
        if denom == 0:
            d = np.nan
        else:
            d = 2.0 * np.logical_and(p, t).sum() / denom
        out[f"dice_c{c}"] = float(d) if not np.isnan(d) else np.nan
        if not np.isnan(d):
            vals.append(d)
    out["dice_mean"] = float(np.mean(vals)) if vals else np.nan
    return out


def _surface(mask: np.ndarray) -> np.ndarray:
    if mask.sum() == 0:
        return mask.astype(bool)
    eroded = ndimage.binary_erosion(mask, iterations=1, border_value=0)
    return np.logical_xor(mask, eroded)


def hd95_per_class(pred: np.ndarray, target: np.ndarray, num_classes: int, spacing=None, include_bg: bool = False, empty_penalty: str = "diagonal") -> Dict[str, float]:
    """Compute per-class HD95.

    If one mask is empty and the other is not, assigning NaN and dropping the
    class makes failed segmentations look artificially good. We instead use the
    physical image diagonal as a conservative finite penalty. If both masks are
    empty, the class is not present and is excluded from the mean.
    """
    out = {}
    start = 0 if include_bg else 1
    vals = []
    if spacing is None:
        spacing = (1.0, 1.0, 1.0)
    spacing = tuple(float(x) for x in spacing)
    diag = float(np.sqrt(sum(((s * max(1, n - 1)) ** 2) for s, n in zip(spacing, target.shape))))
    for c in range(start, num_classes):
        p = pred == c
        t = target == c
        p_sum = int(p.sum())
        t_sum = int(t.sum())
        if p_sum == 0 and t_sum == 0:
            hd = np.nan
        elif p_sum == 0 or t_sum == 0:
            hd = diag if empty_penalty == "diagonal" else np.nan
        else:
            ps = _surface(p)
            ts = _surface(t)
            if ps.sum() == 0 or ts.sum() == 0:
                hd = diag if empty_penalty == "diagonal" else np.nan
            else:
                dt_t = ndimage.distance_transform_edt(~ts, sampling=spacing)
                dt_p = ndimage.distance_transform_edt(~ps, sampling=spacing)
                dists = np.concatenate([dt_t[ps], dt_p[ts]])
                hd = np.percentile(dists, 95) if dists.size else np.nan
        out[f"hd95_c{c}"] = float(hd) if not np.isnan(hd) else np.nan
        if not np.isnan(hd):
            vals.append(hd)
    out["hd95_mean"] = float(np.mean(vals)) if vals else np.nan
    return out


def torch_soft_dice_loss(logits: torch.Tensor, target: torch.Tensor, num_classes: int, ignore_index: int | None = None, eps: float = 1e-5):
    probs = torch.softmax(logits, dim=1)
    if target.ndim == logits.ndim:
        onehot = target.float()
    else:
        target_clamped = target.clamp(0, num_classes - 1).long()
        onehot = torch.nn.functional.one_hot(target_clamped, num_classes).permute(0, 4, 1, 2, 3).float()
    if ignore_index is not None:
        mask = target != ignore_index
        probs = probs * mask.unsqueeze(1)
        onehot = onehot * mask.unsqueeze(1)
    dims = tuple(range(2, logits.ndim))
    inter = (probs * onehot).sum(dims)
    denom = probs.sum(dims) + onehot.sum(dims)
    dice = (2 * inter + eps) / (denom + eps)
    return 1.0 - dice[:, 1:].mean()


def entropy_loss(logits: torch.Tensor, eps: float = 1e-8):
    p = torch.softmax(logits, dim=1)
    return -(p * (p + eps).log()).sum(dim=1).mean()


def confidence_and_margin(probs: torch.Tensor):
    vals, inds = probs.topk(k=2, dim=1)
    conf = vals[:, 0]
    margin = vals[:, 0] - vals[:, 1]
    pred = inds[:, 0]
    return conf, margin, pred


def finite_difference_boundary(x: torch.Tensor) -> torch.Tensor:
    # x: [B,C,H,W,D] probabilities or one-hot maps
    dx = torch.abs(x[:, :, 1:] - x[:, :, :-1]).mean(dim=1, keepdim=True)
    dx = torch.nn.functional.pad(dx, (0,0,0,0,1,0))
    dy = torch.abs(x[:, :, :, 1:] - x[:, :, :, :-1]).mean(dim=1, keepdim=True)
    dy = torch.nn.functional.pad(dy, (0,0,1,0,0,0))
    dz = torch.abs(x[:, :, :, :, 1:] - x[:, :, :, :, :-1]).mean(dim=1, keepdim=True)
    dz = torch.nn.functional.pad(dz, (1,0,0,0,0,0))
    return (dx + dy + dz) / 3.0