import torch
from typing import Optional, Callable


def levenshtein_distance(
        pred: torch.LongTensor,
        target: torch.LongTensor,
        pred_extra: Optional[torch.Tensor] = None,
        target_extra: Optional[torch.Tensor] = None,
        extra_fn: Optional[Callable] = None
) -> torch.FloatTensor:
    """
    Overview:
        Levenshtein Distance, i.e. Edit Distance.
    Arguments:
        - pred (:obj:`torch.LongTensor`): The first tensor to calculate the distance, shape: (N1, )  (N1 >= 0).
        - target (:obj:`torch.LongTensor`): The second tensor to calculate the distance, shape: (N2, )  (N2 >= 0).
        - pred_extra (:obj:`Optional[torch.Tensor]`): Extra tensor to calculate the distance, only works when \
            ``extra_fn`` is not ``None``.
        - target_extra (:obj:`Optional[torch.Tensor]`): Extra tensor to calculate the distance, only works when \
            ``extra_fn`` is not ``None``.
        - extra_fn (:obj:`Optional[Callable]`): The distance function for ``pred_extra`` and \
            ``target_extra``. If set to ``None``, this distance will not be considered.
    Returns:
        - distance (:obj:`torch.FloatTensor`): distance(scalar), shape: (1, ).
    """
    assert (isinstance(pred, torch.Tensor) and isinstance(target, torch.Tensor))
    assert (pred.dtype == torch.long and target.dtype == torch.long), '{}\t{}'.format(pred.dtype, target.dtype)
    assert (pred.device == target.device)
    assert (type(pred_extra) == type(target_extra))
    if not extra_fn:
        assert (not pred_extra)
    N1, N2 = pred.shape[0], target.shape[0]
    assert (N1 >= 0 and N2 >= 0)
    if N1 == 0 or N2 == 0:
        distance = max(N1, N2)
    else:
        dp_array = torch.zeros(N1, N2).float()
        if extra_fn:
            if pred[0] == target[0]:
                extra = extra_fn(pred_extra[0], target_extra[0])
            else:
                extra = 1.
            dp_array[0, :] = torch.arange(0, N2) + extra
            dp_array[:, 0] = torch.arange(0, N1) + extra
        else:
            dp_array[0, :] = torch.arange(0, N2)
            dp_array[:, 0] = torch.arange(0, N1)
        for i in range(1, N1):
            for j in range(1, N2):
                if pred[i] == target[j]:
                    if extra_fn:
                        dp_array[i, j] = dp_array[i - 1, j - 1] + extra_fn(pred_extra[i], target_extra[j])
                    else:
                        dp_array[i, j] = dp_array[i - 1, j - 1]
                else:
                    dp_array[i, j] = min(dp_array[i - 1, j] + 1, dp_array[i, j - 1] + 1, dp_array[i - 1, j - 1] + 1)
        distance = dp_array[N1 - 1, N2 - 1]
    return torch.FloatTensor([distance]).to(pred.device)


def hamming_distance(pred: torch.LongTensor, target: torch.LongTensor, weight=1.) -> torch.LongTensor:
    """
    Overview:
        Hamming Distance.
    Arguments:
        - pred (:obj:`torch.LongTensor`): Pred input, boolean vector(0 or 1).
        - target (:obj:`torch.LongTensor`): Target input, boolean vector(0 or 1).
        - weight (:obj:`torch.LongTensor`): Weight to multiply.
    Returns:
        - distance(:obj:`torch.LongTensor`): Distance (scalar), shape (1, ).
    Shapes:
        - pred & target (:obj:`torch.LongTensor`): shape :math:`(B, N)`, \
            while B is the batch size, N is the dimension
    """
    assert (isinstance(pred, torch.Tensor) and isinstance(target, torch.Tensor))
    assert (pred.dtype == torch.long and target.dtype == torch.long)
    assert (pred.device == target.device)
    assert (pred.shape == target.shape)
    return pred.ne(target).sum(dim=1).float().mul_(weight)