from typing import Optional, List
from functools import partial

import torch
from torch.nn.modules.loss import _Loss
from ._functional import focal_loss_with_logits, to_tensor
from .constants import BINARY_MODE, MULTICLASS_MODE, MULTILABEL_MODE

__all__ = ["FocalLoss"]


class FocalLoss(_Loss):
    def __init__(
        self,
        mode: str,
        alpha: Optional[float] = None,
        gamma: Optional[float] = 2.0,
        ignore_index: Optional[int] = None,
        from_logits: bool = True,
        eps: float = 1e-7,
        reduction: Optional[str] = "mean",
        normalized: bool = False,
        reduced_threshold: Optional[float] = None,
        class_weights: Optional[List[float]] = None,
    ):
        """Compute Focal loss

        Args:
            mode: Loss mode 'binary', 'multiclass' or 'multilabel'
            from_logits: If True, assumes input is raw logits
            eps: Small value used for numerical stability when converting probabilities to logits .
            alpha: Prior probability of having positive value in target.
            gamma: Power factor for dampening weight (focal strength).
            ignore_index: If not None, targets may contain values to be ignored.
                Target values equal to ignore_index will be ignored from loss computation.
            normalized: Compute normalized focal loss (https://arxiv.org/pdf/1909.07829.pdf).
            reduced_threshold: Switch to reduced focal loss. Note, when using this mode you
                should use `reduction="sum"`.
            class_weights: List of weights for each class. If not ``None``, the loss for each class
                is multiplied by the corresponding weight. Only supported for multiclass and
                multilabel modes. Weights do not need to be normalized.

        Shape
             - **y_pred** - torch.Tensor of shape (N, C, H, W)
             - **y_true** - torch.Tensor of shape (N, H, W) or (N, C, H, W)

        Reference
            https://github.com/BloodAxe/pytorch-toolbelt

        """
        assert mode in {BINARY_MODE, MULTILABEL_MODE, MULTICLASS_MODE}
        super().__init__()

        if class_weights is not None and mode == BINARY_MODE:
            raise ValueError("class_weights are not supported with mode=binary")

        self.mode = mode
        self.from_logits = from_logits
        self.ignore_index = ignore_index
        self.reduction = reduction
        self.eps = eps

        self.focal_loss_fn = partial(
            focal_loss_with_logits,
            alpha=alpha,
            gamma=gamma,
            reduced_threshold=reduced_threshold,
            reduction=reduction,
            normalized=normalized,
        )
        self.class_weights = (
            to_tensor(class_weights, dtype=torch.float)
            if class_weights is not None
            else None
        )

    def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:

        if not self.from_logits:
            y_pred = torch.clamp(y_pred, self.eps, 1 - self.eps)

            if self.mode in {BINARY_MODE, MULTILABEL_MODE}:
                # inverse sigmoid
                y_pred = torch.log(y_pred / (1 - y_pred))

            elif self.mode == MULTICLASS_MODE:
                # convert softmax probabilities to log-space
                y_pred = torch.log(y_pred)

        if self.mode == BINARY_MODE:
            y_true = y_true.reshape(-1)
            y_pred = y_pred.reshape(-1)

            if self.ignore_index is not None:
                # Filter predictions with ignore label from loss computation
                not_ignored = y_true != self.ignore_index
                y_pred = y_pred[not_ignored]
                y_true = y_true[not_ignored]

            loss = self.focal_loss_fn(y_pred, y_true)

        elif self.mode in {MULTILABEL_MODE, MULTICLASS_MODE}:
            num_classes = y_pred.size(1)

            # Filter anchors with -1 label from loss computation
            if self.ignore_index is not None:
                not_ignored = y_true != self.ignore_index

            class_losses = []
            for cls in range(num_classes):
                if self.mode == MULTICLASS_MODE:
                    cls_y_true = (y_true == cls).long()
                else:
                    cls_y_true = y_true[:, cls, ...]
                cls_y_pred = y_pred[:, cls, ...]

                if self.ignore_index is not None:
                    cls_y_true = cls_y_true[not_ignored]
                    cls_y_pred = cls_y_pred[not_ignored]

                class_losses.append(self.focal_loss_fn(cls_y_pred, cls_y_true))
            class_losses = torch.stack(class_losses)  # shape (C,)

            if self.class_weights is not None:
                weights = self.class_weights.to(class_losses.device)
                loss = (class_losses * weights).sum() / weights.sum()
            else:
                loss = class_losses.mean()

        return loss
