import torch


def charbonnier_loss(pred: torch.Tensor,
                     target: torch.Tensor,
                     q: float = 0.2,
                     eps: float = 0.01) -> torch.Tensor:
    """Generalized Charbonnier loss function between output and ground truth.

    The loss function is
    .. math::
      loss = ((u-u_gt)^2+(v-v_gt)^2+eps)^q

    Generalized Charbonnier loss was used in LiteFlowNet when fine tuning,
    with eps=0.01 q=0.2.

    Args:
        pred (torch.Tensor): output flow map from flow_estimator
            shape(B, 2, H, W).
        target (torch.Tensor): ground truth flow map shape(B, 2, H, W).
        q (float): the exponent in charbonnier loss.
        eps (float): small constant to numerical stability when
            fine-tuning model. Defaults to 0.01.

    Returns:
        Tensor: loss map with the shape (B, H, W).
    """

    assert pred.shape == target.shape, \
        (f'pred shape {pred.shape} does not match target '
         f'shape {target.shape}.')

    diff = torch.add(pred, -target)

    loss_map = (torch.sum(diff * diff, dim=1) + eps)**q  # shape (B, H, W).

    return loss_map