# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved

from typing import Any, List

import torch

from detectron2.config import CfgNode
from detectron2.structures import Instances

from .mask import MaskLoss
from .segm import SegmentationLoss


class MaskOrSegmentationLoss:
    """
    Mask or segmentation loss as cross-entropy for raw unnormalized scores
    given ground truth labels. Ground truth labels are either defined by coarse
    segmentation annotation, or by mask annotation, depending on the config
    value MODEL.ROI_DENSEPOSE_HEAD.COARSE_SEGM_TRAINED_BY_MASKS
    """

    def __init__(self, cfg: CfgNode):
        """
        Initialize segmentation loss from configuration options

        Args:
            cfg (CfgNode): configuration options
        """
        self.segm_trained_by_masks = (
            cfg.MODEL.ROI_DENSEPOSE_HEAD.COARSE_SEGM_TRAINED_BY_MASKS
        )
        if self.segm_trained_by_masks:
            self.mask_loss = MaskLoss()
        self.segm_loss = SegmentationLoss(cfg)

    def __call__(
        self,
        proposals_with_gt: List[Instances],
        densepose_predictor_outputs: Any,
        packed_annotations: Any,
    ) -> torch.Tensor:
        """
        Compute segmentation loss as cross-entropy between aligned unnormalized
        score estimates and ground truth; with ground truth given
        either by masks, or by coarse segmentation annotations.

        Args:
            proposals_with_gt (list of Instances): detections with associated ground truth data
            densepose_predictor_outputs: an object of a dataclass that contains predictor outputs
                with estimated values; assumed to have the following attributes:
                * coarse_segm - coarse segmentation estimates, tensor of shape [N, D, S, S]
            packed_annotations: packed annotations for efficient loss computation
        Return:
            tensor: loss value as cross-entropy for raw unnormalized scores
                given ground truth labels
        """
        if self.segm_trained_by_masks:
            return self.mask_loss(proposals_with_gt, densepose_predictor_outputs)
        return self.segm_loss(
            proposals_with_gt, densepose_predictor_outputs, packed_annotations
        )

    def fake_value(self, densepose_predictor_outputs: Any) -> torch.Tensor:
        """
        Fake segmentation loss used when no suitable ground truth data
        was found in a batch. The loss has a value 0 and is primarily used to
        construct the computation graph, so that `DistributedDataParallel`
        has similar graphs on all GPUs and can perform reduction properly.

        Args:
            densepose_predictor_outputs: DensePose predictor outputs, an object
                of a dataclass that is assumed to have `coarse_segm`
                attribute
        Return:
            Zero value loss with proper computation graph
        """
        return densepose_predictor_outputs.coarse_segm.sum() * 0