# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/05_vision_metrics.ipynb.

# %% auto #0
__all__ = ['calculate_dsc', 'calculate_haus', 'binary_dice_score', 'multi_dice_score', 'binary_hausdorff_distance',
           'multi_hausdorff_distance', 'calculate_confusion_metrics', 'binary_sensitivity', 'multi_sensitivity',
           'binary_precision', 'multi_precision', 'calculate_lesion_detection_rate', 'binary_lesion_detection_rate',
           'multi_lesion_detection_rate', 'calculate_signed_rve', 'binary_signed_rve', 'multi_signed_rve',
           'AccumulatedDice', 'AccumulatedMultiDice', 'EMACheckpoint']

# %% ../nbs/05_vision_metrics.ipynb #8b6a83ac
import torch
import numpy as np
from monai.metrics import compute_hausdorff_distance, compute_dice, get_confusion_matrix, compute_confusion_matrix_metric
from scipy.ndimage import label as scipy_label
from .vision_data import pred_to_binary_mask, batch_pred_to_multiclass_mask
from fastai.learner import Metric
from fastai.callback.tracker import TrackerCallback

# %% ../nbs/05_vision_metrics.ipynb #5c16dc6c-e07a-44b5-85af-79bd6c4ce390
def calculate_dsc(pred: torch.Tensor, targ: torch.Tensor) -> torch.Tensor:
    """Calculate Dice score using MONAI's compute_dice.

    Accepts tensors of various shapes and automatically reshapes to 5D format.

    Args:
        pred: Binary prediction tensor. Accepts:
              - [D, H, W] single 3D volume
              - [C, D, H, W] single volume with channel
              - [B, C, D, H, W] batched volumes
        targ: Binary target tensor (same shape options as pred).

    Returns:
        Dice score(s). Single value for 3D/4D input, tensor of values for 5D batch.
    """
    # Normalize to 5D: [B, C, D, H, W]
    if pred.ndim == 3:  # [D, H, W] -> [1, 1, D, H, W]
        pred = pred.unsqueeze(0).unsqueeze(0)
        targ = targ.unsqueeze(0).unsqueeze(0)
    elif pred.ndim == 4:  # [C, D, H, W] -> [1, C, D, H, W]
        pred = pred.unsqueeze(0)
        targ = targ.unsqueeze(0)

    return torch.Tensor([compute_dice(p[None], t[None]) for p, t in zip(pred, targ)])

# %% ../nbs/05_vision_metrics.ipynb #4f815b1d-ea53-4f65-a7f0-299ea54e872b
def calculate_haus(pred: torch.Tensor, targ: torch.Tensor) -> torch.Tensor:
    """Compute 95th percentile Hausdorff distance (HD95) using MONAI.

    HD95 is more robust than standard Hausdorff distance as it ignores
    the top 5% of outlier distances.

    Accepts tensors of various shapes and automatically reshapes to 5D format.

    Args:
        pred: Binary prediction tensor. Accepts:
              - [D, H, W] single 3D volume
              - [C, D, H, W] single volume with channel
              - [B, C, D, H, W] batched volumes
        targ: Binary target tensor (same shape options as pred).

    Returns:
        HD95 value(s). Single value for 3D/4D input, tensor of values for 5D batch.
    """
    # Normalize to 5D: [B, C, D, H, W]
    if pred.ndim == 3:  # [D, H, W] -> [1, 1, D, H, W]
        pred = pred.unsqueeze(0).unsqueeze(0)
        targ = targ.unsqueeze(0).unsqueeze(0)
    elif pred.ndim == 4:  # [C, D, H, W] -> [1, C, D, H, W]
        pred = pred.unsqueeze(0)
        targ = targ.unsqueeze(0)

    return torch.Tensor([compute_hausdorff_distance(p[None], t[None], percentile=95) for p, t in zip(pred, targ)])

# %% ../nbs/05_vision_metrics.ipynb #430d64f8-8cd9-4a88-ad20-5f73bebbf12f
def binary_dice_score(act: torch.Tensor, targ: torch.Tensor) -> torch.Tensor:
    """Calculates the mean Dice score for binary semantic segmentation tasks.
    
    Args:
        act: Activation tensor with dimensions [B, C, W, H, D].
        targ: Target masks with dimensions [B, C, W, H, D].

    Returns:
        Mean Dice score.
    """
    pred = pred_to_binary_mask(act)
    dsc = calculate_dsc(pred.cpu(), targ.cpu())

    return torch.mean(dsc)

# %% ../nbs/05_vision_metrics.ipynb #48ba4382-eeb0-46d7-8f84-515313c7c27c
def multi_dice_score(act: torch.Tensor, targ: torch.Tensor) -> torch.Tensor:
    """Calculate the mean Dice score for each class in multi-class semantic 
    segmentation tasks.

    Args:
        act: Activation tensor with dimensions [B, C, W, H, D].
        targ: Target masks with dimensions [B, C, W, H, D].

    Returns:
        Mean Dice score for each class.
    """
    pred, n_classes = batch_pred_to_multiclass_mask(act)
    binary_dice_scores = []

    for c in range(1, n_classes):
        c_pred, c_targ = torch.where(pred == c, 1, 0), torch.where(targ == c, 1, 0)
        dsc = calculate_dsc(c_pred, c_targ)
        binary_dice_scores.append(float(torch.nanmean(dsc)))

    return torch.Tensor(binary_dice_scores)

# %% ../nbs/05_vision_metrics.ipynb #a390762f-d1a9-4674-b099-2369769f4198
def binary_hausdorff_distance(act: torch.Tensor, targ: torch.Tensor) -> torch.Tensor:
    """Calculate the mean HD95 for binary semantic segmentation tasks.
    
    Args:
        act: Activation tensor with dimensions [B, C, W, H, D].
        targ: Target masks with dimensions [B, C, W, H, D].

    Returns:
        Mean HD95.
    """
    pred = pred_to_binary_mask(act)
    haus = calculate_haus(pred.cpu(), targ.cpu())
    return torch.mean(haus)

# %% ../nbs/05_vision_metrics.ipynb #ea94dea5
def multi_hausdorff_distance(act: torch.Tensor, targ: torch.Tensor) -> torch.Tensor:
    """Calculate the mean HD95 for each class in multi-class semantic segmentation tasks.
    
    Args:
        act: Activation tensor with dimensions [B, C, W, H, D].
        targ: Target masks with dimensions [B, C, W, H, D].

    Returns:
        Mean HD95 for each class.
    """
    pred, n_classes = batch_pred_to_multiclass_mask(act)
    binary_haus = []

    for c in range(1, n_classes):
        c_pred, c_targ = torch.where(pred==c, 1, 0), torch.where(targ==c, 1, 0)
        haus = calculate_haus(c_pred, c_targ)
        binary_haus.append(float(torch.nanmean(haus)))
    return torch.Tensor(binary_haus)

# %% ../nbs/05_vision_metrics.ipynb #fdkdak1uf9s
def calculate_confusion_metrics(pred: torch.Tensor, targ: torch.Tensor, metric_name: str) -> torch.Tensor:
    """Calculate confusion matrix-based metric using MONAI.
    
    Args:
        pred: Binary prediction tensor [B, C, W, H, D].
        targ: Binary target tensor [B, C, W, H, D].
        metric_name: One of "sensitivity", "precision", "specificity", "f1 score".
    
    Returns:
        Metric values for each sample in batch.
    """
    # get_confusion_matrix expects one-hot format and returns [B, n_class, 4] where 4 = [TP, FP, TN, FN]
    confusion_matrix = get_confusion_matrix(pred, targ, include_background=False)
    metric = compute_confusion_matrix_metric(metric_name, confusion_matrix)
    return metric

# %% ../nbs/05_vision_metrics.ipynb #zbxmbdwq0ao
def binary_sensitivity(act: torch.Tensor, targ: torch.Tensor) -> torch.Tensor:
    """Calculate mean sensitivity (recall) for binary segmentation.
    
    Sensitivity = TP / (TP + FN) - measures the proportion of actual positives
    that are correctly identified.
    
    Args:
        act: Activation tensor [B, C, W, H, D].
        targ: Target masks [B, C, W, H, D].
    
    Returns:
        Mean sensitivity score.
    """
    pred = pred_to_binary_mask(act)
    sens = calculate_confusion_metrics(pred.cpu(), targ.cpu(), "sensitivity")
    return torch.nanmean(sens)

# %% ../nbs/05_vision_metrics.ipynb #q0f40h69zan
def multi_sensitivity(act: torch.Tensor, targ: torch.Tensor) -> torch.Tensor:
    """Calculate mean sensitivity for each class in multi-class segmentation.
    
    Args:
        act: Activation tensor [B, C, W, H, D].
        targ: Target masks [B, C, W, H, D].
    
    Returns:
        Mean sensitivity for each class.
    """
    pred, n_classes = batch_pred_to_multiclass_mask(act)
    class_sens = []
    
    for c in range(1, n_classes):
        c_pred = torch.where(pred == c, 1, 0)
        c_targ = torch.where(targ == c, 1, 0)
        sens = calculate_confusion_metrics(c_pred, c_targ, "sensitivity")
        class_sens.append(float(torch.nanmean(sens)))
    
    return torch.Tensor(class_sens)

# %% ../nbs/05_vision_metrics.ipynb #dlnx9vkx98j
def binary_precision(act: torch.Tensor, targ: torch.Tensor) -> torch.Tensor:
    """Calculate mean precision for binary segmentation.
    
    Precision = TP / (TP + FP) - measures the proportion of positive predictions
    that are actually correct.
    
    Args:
        act: Activation tensor [B, C, W, H, D].
        targ: Target masks [B, C, W, H, D].
    
    Returns:
        Mean precision score.
    """
    pred = pred_to_binary_mask(act)
    prec = calculate_confusion_metrics(pred.cpu(), targ.cpu(), "precision")
    return torch.nanmean(prec)

# %% ../nbs/05_vision_metrics.ipynb #cwb5sivj8es
def multi_precision(act: torch.Tensor, targ: torch.Tensor) -> torch.Tensor:
    """Calculate mean precision for each class in multi-class segmentation.
    
    Args:
        act: Activation tensor [B, C, W, H, D].
        targ: Target masks [B, C, W, H, D].
    
    Returns:
        Mean precision for each class.
    """
    pred, n_classes = batch_pred_to_multiclass_mask(act)
    class_prec = []
    
    for c in range(1, n_classes):
        c_pred = torch.where(pred == c, 1, 0)
        c_targ = torch.where(targ == c, 1, 0)
        prec = calculate_confusion_metrics(c_pred, c_targ, "precision")
        class_prec.append(float(torch.nanmean(prec)))
    
    return torch.Tensor(class_prec)

# %% ../nbs/05_vision_metrics.ipynb #flomm5fkz69
def calculate_lesion_detection_rate(pred: torch.Tensor, targ: torch.Tensor, threshold: float = 0.0) -> torch.Tensor:
    """Calculate lesion-wise detection rate.

    For each connected component (lesion) in the target, check if it is
    detected by the prediction. Detection criteria depends on threshold:
    - threshold=0: any overlap counts as detected
    - threshold>0: per-lesion Dice score must exceed threshold

    Args:
        pred: Binary prediction tensor [B, C, W, H, D].
        targ: Binary target tensor [B, C, W, H, D].
        threshold: Minimum Dice score for a lesion to be considered detected.
                   Default 0.0 means any overlap counts as detected.

    Returns:
        Detection rate (detected lesions / total lesions) for each sample.
    """
    detection_rates = []

    for p, t in zip(pred, targ):
        p_np = p.squeeze().cpu().numpy()
        t_np = t.squeeze().cpu().numpy()

        # Label connected components in target
        labeled_targ, n_lesions = scipy_label(t_np)

        if n_lesions == 0:
            detection_rates.append(float('nan'))
            continue

        detected = 0
        for lesion_id in range(1, n_lesions + 1):
            lesion_mask = (labeled_targ == lesion_id)

            if threshold == 0.0:
                # Original behavior: any overlap counts as detected
                overlap = (p_np * lesion_mask).sum()
                if overlap > 0:
                    detected += 1
            else:
                # Compute per-lesion Dice score
                pred_in_lesion = p_np * lesion_mask
                intersection = (pred_in_lesion * lesion_mask).sum()
                lesion_vol = lesion_mask.sum()
                pred_vol = pred_in_lesion.sum()

                if (lesion_vol + pred_vol) > 0:
                    dice = 2 * intersection / (lesion_vol + pred_vol)
                    if dice > threshold:
                        detected += 1

        detection_rates.append(detected / n_lesions)

    return torch.Tensor(detection_rates)

# %% ../nbs/05_vision_metrics.ipynb #jyjuvh9f09n
def binary_lesion_detection_rate(act: torch.Tensor, targ: torch.Tensor, threshold: float = 0.0) -> torch.Tensor:
    """Calculate mean lesion detection rate for binary segmentation.

    Args:
        act: Activation tensor [B, C, W, H, D].
        targ: Target masks [B, C, W, H, D].
        threshold: Minimum Dice score for a lesion to be considered detected.
                   Default 0.0 means any overlap counts as detected.

    Returns:
        Mean lesion detection rate.
    """
    pred = pred_to_binary_mask(act)
    ldr = calculate_lesion_detection_rate(pred.cpu(), targ.cpu(), threshold)
    return torch.nanmean(ldr)

# %% ../nbs/05_vision_metrics.ipynb #a5cwiazhpp
def multi_lesion_detection_rate(act: torch.Tensor, targ: torch.Tensor, threshold: float = 0.0) -> torch.Tensor:
    """Calculate mean lesion detection rate for each class in multi-class segmentation.

    Args:
        act: Activation tensor [B, C, W, H, D].
        targ: Target masks [B, C, W, H, D].
        threshold: Minimum Dice score for a lesion to be considered detected.
                   Default 0.0 means any overlap counts as detected.

    Returns:
        Mean lesion detection rate for each class.
    """
    pred, n_classes = batch_pred_to_multiclass_mask(act)
    class_ldr = []

    for c in range(1, n_classes):
        c_pred = torch.where(pred == c, 1, 0)
        c_targ = torch.where(targ == c, 1, 0)
        ldr = calculate_lesion_detection_rate(c_pred, c_targ, threshold)
        class_ldr.append(float(torch.nanmean(ldr)))

    return torch.Tensor(class_ldr)

# %% ../nbs/05_vision_metrics.ipynb #8aat811h3zo
def calculate_signed_rve(pred: torch.Tensor, targ: torch.Tensor) -> torch.Tensor:
    """Calculate signed Relative Volume Error.
    
    RVE = (pred_volume - targ_volume) / targ_volume
    
    Positive values indicate over-segmentation (model predicts too large),
    negative values indicate under-segmentation (model predicts too small).
    
    Args:
        pred: Binary prediction tensor [B, C, W, H, D].
        targ: Binary target tensor [B, C, W, H, D].
    
    Returns:
        Signed RVE for each sample in batch.
    """
    rve_values = []
    
    for p, t in zip(pred, targ):
        pred_vol = p.sum().float()
        targ_vol = t.sum().float()
        
        if targ_vol == 0:
            rve_values.append(float('nan'))
        else:
            rve = (pred_vol - targ_vol) / targ_vol
            rve_values.append(rve.item())
    
    return torch.Tensor(rve_values)

# %% ../nbs/05_vision_metrics.ipynb #a18cg6z9qtr
def binary_signed_rve(act: torch.Tensor, targ: torch.Tensor) -> torch.Tensor:
    """Calculate mean signed RVE for binary segmentation.
    
    Args:
        act: Activation tensor [B, C, W, H, D].
        targ: Target masks [B, C, W, H, D].
    
    Returns:
        Mean signed RVE.
    """
    pred = pred_to_binary_mask(act)
    rve = calculate_signed_rve(pred.cpu(), targ.cpu())
    return torch.nanmean(rve)

# %% ../nbs/05_vision_metrics.ipynb #aetwom9qd3p
def multi_signed_rve(act: torch.Tensor, targ: torch.Tensor) -> torch.Tensor:
    """Calculate mean signed RVE for each class in multi-class segmentation.
    
    Args:
        act: Activation tensor [B, C, W, H, D].
        targ: Target masks [B, C, W, H, D].
    
    Returns:
        Mean signed RVE for each class.
    """
    pred, n_classes = batch_pred_to_multiclass_mask(act)
    class_rve = []
    
    for c in range(1, n_classes):
        c_pred = torch.where(pred == c, 1, 0)
        c_targ = torch.where(targ == c, 1, 0)
        rve = calculate_signed_rve(c_pred, c_targ)
        class_rve.append(float(torch.nanmean(rve)))
    
    return torch.Tensor(class_rve)

# %% ../nbs/05_vision_metrics.ipynb #js7clwse6n
class AccumulatedDice(Metric):
    """nnU-Net-style accumulated Dice metric for reliable pseudo dice during training.

    Instead of averaging per-batch Dice scores, this metric accumulates
    true positives, false positives, and false negatives across ALL validation
    batches, then computes Dice from the totals. This gives more weight to
    batches with more foreground voxels and is more statistically robust.

    Args:
        n_classes: Number of classes including background (default: 2 for binary).
        include_background: Whether to include background in metric (default: False).

    Example:
        ```python
        learn = Learner(dls, model, loss_func=loss_func,
                       metrics=[AccumulatedDice(n_classes=2)])
        
        # For checkpoint selection based on accumulated dice:
        save_best = SaveModelCallback(
            monitor='accumulated_dice',
            comp=np.greater,  # Higher dice is better
            fname='best_model'
        )
        ```
    """
    def __init__(self, n_classes: int = 2, include_background: bool = False):
        self.n_classes = n_classes
        self.include_background = include_background
        self.start_class = 0 if include_background else 1

    def reset(self):
        """Called at start of validation epoch."""
        n_fg_classes = self.n_classes if self.include_background else self.n_classes - 1
        self.tp = torch.zeros(n_fg_classes, dtype=torch.float64)
        self.fp = torch.zeros(n_fg_classes, dtype=torch.float64)
        self.fn = torch.zeros(n_fg_classes, dtype=torch.float64)

    def accumulate(self, learn):
        """Called after each validation batch to accumulate TP/FP/FN."""
        pred = learn.pred  # Model output [B, C, D, H, W]
        targ = learn.y     # Target [B, 1, D, H, W] or [B, D, H, W]

        # Get predicted segmentation
        pred_seg = pred.argmax(dim=1)  # [B, D, H, W]

        # Ensure target has same shape
        if targ.ndim == pred_seg.ndim + 1:
            targ = targ.squeeze(1)

        # Accumulate TP/FP/FN for each foreground class
        idx = 0
        for c in range(self.start_class, self.n_classes):
            c_pred = (pred_seg == c)
            c_targ = (targ == c)

            self.tp[idx] += (c_pred & c_targ).sum().float().cpu()
            self.fp[idx] += (c_pred & ~c_targ).sum().float().cpu()
            self.fn[idx] += (~c_pred & c_targ).sum().float().cpu()
            idx += 1

    @property
    def value(self):
        """Compute Dice from accumulated counts at end of epoch."""
        dice = 2 * self.tp / (2 * self.tp + self.fp + self.fn + 1e-8)
        return dice.mean().item()

    @property
    def name(self):
        return 'accumulated_dice'

# %% ../nbs/05_vision_metrics.ipynb #a9kkhgt8ge8
class AccumulatedMultiDice(AccumulatedDice):
    """Multi-class version of AccumulatedDice that returns per-class Dice scores.

    Instead of returning a single mean Dice, this returns a tensor with the
    Dice score for each foreground class. Useful for monitoring per-class
    performance during training.

    Example:
        ```python
        # For 3-class segmentation (background + 2 foreground classes)
        learn = Learner(dls, model, loss_func=loss_func,
                       metrics=[AccumulatedMultiDice(n_classes=3)])
        ```
    """
    @property
    def value(self):
        """Return per-class Dice scores."""
        dice = 2 * self.tp / (2 * self.tp + self.fp + self.fn + 1e-8)
        return dice  # Returns tensor, fastai will display all values

    @property
    def name(self):
        return 'accumulated_multi_dice'

# %% ../nbs/05_vision_metrics.ipynb #mwwwrlj5sm
class EMACheckpoint(TrackerCallback):
    """Save model checkpoint based on EMA of a monitored metric (nnU-Net style).

    Instead of saving the best model based on a single (noisy) epoch metric,
    this tracks the exponential moving average and saves when the EMA improves.
    More robust for patch-based training where per-epoch metrics fluctuate.

    Formula: ema = momentum * previous_ema + (1 - momentum) * current_value

    Unlike SaveModelCallback, this does NOT auto-load the best model after
    training. Load explicitly with ``learn.load(fname)``.

    Args:
        monitor: Metric name to track (default: 'accumulated_dice').
        momentum: EMA momentum (default: 0.9, matching nnU-Net).
            Higher momentum = more smoothing. Range: (0, 1).
            nnU-Net uses 0.9 (keeps 90% of history, adds 10% of current epoch).
        comp: Comparison function (default: np.greater for higher-is-better).
        fname: Model save filename (default: 'best_model').
        with_opt: Whether to save optimizer state (default: False).

    Example:
        ```python
        save_best = EMACheckpoint(
            monitor='accumulated_dice',
            momentum=0.9,
            fname='best_model'
        )
        learn.fit_one_cycle(30, lr, cbs=[save_best])

        # Load best model after training:
        learn.load('best_model')

        # Access EMA history for plotting:
        save_best.ema_history
        ```
    """
    order = 60  # Same priority as SaveModelCallback

    def __init__(self, monitor='accumulated_dice', momentum=0.9, comp=np.greater,
                 fname='best_model', with_opt=False):
        super().__init__(monitor=monitor, comp=comp)
        self.fname = fname
        self.with_opt = with_opt
        self.momentum = momentum

    def before_fit(self):
        super().before_fit()  # Sets self.idx, self.best via TrackerCallback
        self.ema_value = None
        self.ema_history = []

    def after_epoch(self):
        val = self.recorder.values[-1][self.idx]
        if isinstance(val, torch.Tensor): val = val.item()

        self.ema_value = val if self.ema_value is None else (
            self.momentum * self.ema_value + (1 - self.momentum) * val)
        self.ema_history.append(self.ema_value)

        if self.comp(self.ema_value, self.best):
            self.best = self.ema_value
            self.learn.save(self.fname, with_opt=self.with_opt)
