import torch

import focal_loss_cuda


class FocalLoss(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx,
        cls_output,
        cls_targets_at_level,
        num_positives_sum,
        num_real_classes,
        alpha,
        gamma,
        label_smoothing=0.0,
    ):
        loss, partial_grad = focal_loss_cuda.forward(
            cls_output,
            cls_targets_at_level,
            num_positives_sum,
            num_real_classes,
            alpha,
            gamma,
            label_smoothing,
        )

        ctx.save_for_backward(partial_grad, num_positives_sum)
        return loss

    @staticmethod
    def backward(ctx, grad_loss):
        partial_grad, num_positives_sum = ctx.saved_tensors

        # The backward kernel is actually in-place to save memory space,
        # partial_grad and grad_input are the same tensor.
        grad_input = focal_loss_cuda.backward(grad_loss, partial_grad, num_positives_sum)

        return grad_input, None, None, None, None, None, None


def focal_loss(
    cls_output: torch.Tensor,
    cls_targets_at_level: torch.Tensor,
    num_positive_sum: torch.Tensor,
    num_real_classes: int,
    alpha: float,
    gamma: float,
    label_smoothing: float = 0.0,
) -> torch.Tensor:
    """Fused focal loss function."""
    return FocalLoss.apply(
        cls_output,
        cls_targets_at_level,
        num_positive_sum,
        num_real_classes,
        alpha,
        gamma,
        label_smoothing,
    )