import torch
from typing import Callable


def chamfer(input: torch.Tensor, target: torch.Tensor, dt2_fcn: Callable[[torch.Tensor], torch.Tensor],
            ras_fcn: Callable[[torch.Tensor], torch.Tensor], symmetric: bool = False,
            reduction: str = 'mean') -> torch.Tensor:
    """
    Generic Chamfer distance function.

    This chamfer distance implementation is designed to be flexible and can work with arbitrary data (as long as it can
    be represented as a tensor). The approach is as follows:
        1. the squared total distance transform is computed by applying the provided `dt2_fcn` to the input tensor.
        2. the total distance transform is computed by square-rooting the squared total distance transform.
        3. the target mask image is computed by applying the `ras_fcn` to the target tensor.
        4. the chamfer tensor is computed by computing the hadamard product of the total distance transform and target
           mask.
        5. the chamfer tensor is reduced using the specified reduction.

    Examples:
        - if the input represented line segments, then the `dt2_fcn` would need to compute the total (min across each
          segment) squared distance transform.
        - if the target represented line segments, then the `ras_fcn` would need to compute a binary raster of those
          segments (need not be differentiable)
        - if the target represented a binary target image, then the `ras_fcn` would just be None.

    Args:
        input: batch of input data
        target: batch of target data
        dt2_fcn: function to transform input into a batch of squared distance transform images
        ras_fcn: function to transform target into a batch of rasters; can be None if target is already a batch of
          rasters
        symmetric: if True, then the symmetric distance is computed (default False)
        reduction (string, optional): Specifies the reduction to apply to the output:
            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,
            ``'mean'``: the sum of the output will be divided by the number of
            elements in the output, ``'sum'``: the output will be summed. Default: ``'mean'``

    Returns:
        Tensor containing the chamfer distance. For sum of mean reduction this will be scalar; for none it will be
        [batch, *] where * represents the shape of the grid.

    """
    if dt2_fcn is None:
        ic = input.sqrt()
    else:
        ic = dt2_fcn(input).sqrt()

    if ras_fcn is None:
        tc = target
    else:
        tc = ras_fcn(target)

    cham = ic * tc

    if symmetric:
        ic2 = ras_fcn(input)
        tc2 = dt2_fcn(target).sqrt()

        cham2 = ic2 * tc2
        cham += cham2

    if reduction == 'sum':
        return cham.sum()
    if reduction == 'mean':
        return cham.mean()

    return cham