from typing import Literal, Any

import numpy as np
import torch
from scipy.ndimage import distance_transform_edt as distance
from skimage.segmentation import find_boundaries
from torch import Tensor
from FastGeodis import generalised_geodesic2d


def compute_sdf(binary_mask: Any, normalize=True) -> Tensor:
    """Compute the signed distance map of binary mask.

    :param input: segmentation, shape = (num_channels, x, y) or (x, y)
    :param output: the Signed Distance Map (SDM)
    :return: the Signed Distance Map (SDM). Result currently normalized sdf to [-1,1].

    :Notes: Inspired by source https://github.com/JunMa11/SegWithDistMap/blob/master/code/train_LA_MultiHead_SDF_L1.py .

    """
    binary_mask = np.asarray(binary_mask).astype(np.float64)

    ## if not binary mask, really makes trouble so raise error if not just 0 and 1
    if not np.all(np.isin(binary_mask, [0, 1])):
        raise ValueError(
            f"Binary mask should only contain 0 and 1. Found {np.unique(binary_mask)}"
        )

    sdf = np.zeros(binary_mask.shape)

    if len(binary_mask.shape) > 2:
        for b in range(binary_mask.shape[0]):
            sdf[b] = compute_sdf(binary_mask[b], normalize=normalize)
        ## convert to tensor if needed
        return torch.tensor(sdf, dtype=torch.float32)

    if np.all(binary_mask == 0):
        ## If empty mask, all pixels are outside (negative SDF)
        return -torch.ones(binary_mask.shape, dtype=torch.float32)
    elif np.all(binary_mask == 1):
        ## If full mask, all pixels are inside (positive SDF)
        return torch.ones(binary_mask.shape, dtype=torch.float32)

    negmask = 1 - binary_mask

    # ========== DTYPE FIX START ==========
    # - torch.tensor(distance(...)) was creating float64 tensors by default
    # + explicitly convert to float32 to match pred dtype and avoid "Found dtype Double but expected Float" error
    dist_to_background = torch.tensor(distance(binary_mask), dtype=torch.float32)
    dist_to_foreground = torch.tensor(distance(negmask), dtype=torch.float32)
    # ========== DTYPE FIX END ==========

    if normalize:
        dist_to_foreground_normalized = dist_to_foreground / torch.max(
            dist_to_foreground
        )
        dist_to_background_normalized = dist_to_background / torch.max(
            dist_to_background
        )
        ## Do normalization to [-1, 1] here - positive inside, negative outside
        sdf = dist_to_background_normalized - dist_to_foreground_normalized
    else:
        sdf = dist_to_background - dist_to_foreground

    ## Set boundary to 0 (usually, it is set to -1)
    boundary = torch.tensor(
        find_boundaries(binary_mask, mode="inner"), dtype=torch.float32
    )
    sdf[boundary == 1] = 0
    return sdf


def _compute_sdf_fast_single(
    x: Tensor,
    v: float,
    lamb: float,
    iterations: int,
    normalize: bool,
    clip_negatives=False,
) -> Tensor:
    """Compute SDF for a single (1, 1, H, W) tensor."""
    # Handle edge cases: all background or all foreground
    if x.max() == 0:
        # All background - return negative (or zero if clipping)
        if clip_negatives:
            return torch.zeros_like(x)
        elif normalize:
            return -torch.ones_like(x)
        else:
            return -torch.ones_like(x) * (x.size(2) ** 2 + x.size(3) ** 2) ** 0.5
    if x.min() == 1:
        # All foreground - return positive
        if normalize:
            return torch.ones_like(x)
        else:
            return torch.ones_like(x) * (x.size(2) ** 2 + x.size(3) ** 2) ** 0.5

    # FastGeodis computes distance FROM the zero-valued regions of the mask
    # dist_from_bg: 0 at background, increases into foreground
    dist_from_bg = generalised_geodesic2d(x, x, v, lamb, iterations)

    # dist_from_fg: 0 at foreground, increases into background
    dist_from_fg = generalised_geodesic2d(1 - x, 1 - x, v, lamb, iterations)

    if normalize:
        # Safe normalization per-sample: avoid division by zero
        max_from_bg = dist_from_bg.max()
        max_from_fg = dist_from_fg.max()

        if max_from_bg > 0:
            dist_from_bg = dist_from_bg / max_from_bg
        if max_from_fg > 0:
            dist_from_fg = dist_from_fg / max_from_fg

    # SDF convention: positive inside (foreground), negative outside (background)
    sdf = dist_from_bg - dist_from_fg

    if clip_negatives:
        sdf = torch.clamp(sdf, min=0)

    return sdf


def compute_sdf_fast(
    binary_mask: Tensor, normalize: bool = True, clip_negatives: bool = False
) -> Tensor:
    """Compute SDF using FastGeodis for GPU-compatible computation.

    This is a GPU-compatible alternative to compute_sdf that uses FastGeodis
    for distance computation, avoiding CPU transfers during training.

    Args:
        binary_mask: Input tensor of shape (B, C, H, W) with values in [0, 1].
        normalize: If True, normalize SDF to [-1, 1] range per sample.
        clip_negatives: If True, clip negative values to 0 (only keep interior distances).

    Returns:
        SDF tensor of same shape as input (B, C, H, W).
        Positive inside (foreground), negative outside (background).
        If clip_negatives=True, all values are >= 0.
    """
    assert binary_mask.max() <= 1.0, f"Max value {binary_mask.max()} > 1.0"
    assert binary_mask.min() >= 0.0, f"Min value {binary_mask.min()} < 0.0"

    x = binary_mask.float()
    B, C, H, W = x.shape

    # FastGeodis parameters
    v = 1e10
    lamb = 0.0  # 0.0 = Euclidean, 1.0 = Geodesic
    iterations = 2

    # FastGeodis only supports single batch, so we loop over batch and channel
    sdf_results = []
    for b in range(B):
        channel_results = []
        for c in range(C):
            # Extract single sample as (1, 1, H, W)
            single = x[b : b + 1, c : c + 1, :, :]
            sdf_single = _compute_sdf_fast_single(
                single, v, lamb, iterations, normalize, clip_negatives
            )
            channel_results.append(sdf_single)
        # Stack channels: (C, H, W)
        sdf_results.append(torch.cat(channel_results, dim=1))

    # Stack batches: (B, C, H, W)
    return torch.cat(sdf_results, dim=0)


def sdf_difference_weighted_bceloss_from_logits(
    pred,
    target,
    reduction: Literal["mean", "sum", "none"] = "mean",
    normalize: bool = True,
    clip_negatives=False,
    sdf_weights: Tensor | None = None,
):
    """Compute the weighted binary cross entropy loss with respect to the SDF.

    Args:
        pred: Prediction logits (B, C, H, W).
        target: Target mask (B, C, H, W).
        reduction: Reduction mode ('mean', 'sum', 'none').
        normalize: If True, normalize SDFs to [-1, 1] range.
        clip_negatives: If True, clip negative SDF values to 0.
        sdf_weights: Precomputed SDF weights. If None, will be computed.
            Pass this to avoid redundant SDF computation when using multiple
            SDF-weighted losses together.

    Returns:
        Weighted BCE loss.
    """
    bce_map = torch.nn.functional.binary_cross_entropy_with_logits(
        pred, target, reduction="none"
    )
    if sdf_weights is None:
        sdf_weights = sdf_difference_map(
            pred > 0, target, normalize=normalize, clip_negatives=clip_negatives
        )
    if reduction == "mean":
        return torch.mean(bce_map * sdf_weights)
    elif reduction == "sum":
        return torch.sum(bce_map * sdf_weights)
    elif reduction == "none":
        return bce_map * sdf_weights


def sdf_difference_map(mask1, mask2, normalize=True, clip_negatives=False):
    """Compute the absolute difference between SDFs of two masks.

    Args:
        mask1: First mask tensor (B, C, H, W).
        mask2: Second mask tensor (B, C, H, W).
        normalize: If True, normalize SDFs to [-1, 1] range.
        clip_negatives: If True, clip negative SDF values to 0.

    Returns:
        Absolute difference of SDFs, same shape as input.
        Output is on the same device as input (no unnecessary transfers).
    """
    # Binarize masks to handle both binary and probability inputs
    mask1_binary = (mask1 > 0.5).float()
    mask2_binary = (mask2 > 0.5).float()
    sdf1 = compute_sdf_fast(
        mask1_binary, normalize=normalize, clip_negatives=clip_negatives
    )
    sdf2 = compute_sdf_fast(
        mask2_binary, normalize=normalize, clip_negatives=clip_negatives
    )
    # No .to(device) needed - compute_sdf_fast preserves input device
    return torch.abs(sdf1 - sdf2)


if __name__ == "__main__":
    from torchvision.transforms.functional import resize
    import matplotlib.pyplot as plt
    from FastGeodis import generalised_geodesic2d

    x = (
        resize(torch.randint(0, 2, (1, 1, 5, 5), dtype=torch.float32), [480, 480]) > 0.5
    ) * 1.0

    device = "cuda" if torch.cuda.is_available() else "cpu"

    mask_pt = x  # torch.ones_like(x)
    # mask_pt[..., 100, 100] = x

    v = 1e10
    # lamb = 0.0 (Euclidean) or 1.0 (Geodesic) or (0.0, 1.0) (mixture)
    lamb = 0.0
    iterations = 2
    euclidean_dist_neg = generalised_geodesic2d(x, x, v, lamb, iterations)
    euclidean_dist_pos = generalised_geodesic2d(1 - x, 1 - x, v, lamb, iterations)

    norm_sdf_pos = euclidean_dist_pos / euclidean_dist_pos.max()
    norm_sdf_neg = euclidean_dist_neg / euclidean_dist_neg.max()

    unnormalized_sdf_map = euclidean_dist_pos[0, 0, ...] - euclidean_dist_neg[0, 0, ...]
    sdf_map = norm_sdf_pos - norm_sdf_neg

    fig, axes = plt.subplots(1, 3, sharex=True, sharey=True)
    axes = axes.flat

    p1 = axes[0].imshow(x[0, 0, ...])
    p2 = axes[1].imshow(unnormalized_sdf_map)
    p3 = axes[2].imshow(sdf_map[0, 0, ...])

    plt.colorbar(p1, ax=axes[0])
    plt.colorbar(p2, ax=axes[1])
    plt.colorbar(p3, ax=axes[2])
    plt.show()

    # euclidean_dist = np.squeeze(euclidean_dist.cpu().numpy())
