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

from typing import Any
import torch
from torch.nn import functional as F

from detectron2.config import CfgNode
from detectron2.layers import ConvTranspose2d

from densepose.modeling.confidence import DensePoseConfidenceModelConfig
from densepose.modeling.utils import initialize_module_params
from densepose.structures import decorate_cse_predictor_output_class_with_confidences


class DensePoseEmbeddingConfidencePredictorMixin:
    """
    Predictor contains the last layers of a DensePose model that take DensePose head
    outputs as an input and produce model outputs. Confidence predictor mixin is used
    to generate confidences for coarse segmentation estimated by some
    base predictor. Several assumptions need to hold for the base predictor:
    1) the `forward` method must return CSE DensePose head outputs,
        tensor of shape [N, D, H, W]
    2) `interp2d` method must be defined to perform bilinear interpolation;
        the same method is typically used for masks and confidences
    Confidence predictor mixin provides confidence estimates, as described in:
        N. Neverova et al., Correlated Uncertainty for Learning Dense Correspondences
            from Noisy Labels, NeurIPS 2019
        A. Sanakoyeu et al., Transferring Dense Pose to Proximal Animal Classes, CVPR 2020
    """

    def __init__(self, cfg: CfgNode, input_channels: int):
        """
        Initialize confidence predictor using configuration options.

        Args:
            cfg (CfgNode): configuration options
            input_channels (int): number of input channels
        """
        # we rely on base predictor to call nn.Module.__init__
        super().__init__(cfg, input_channels)  # pyre-ignore[19]
        self.confidence_model_cfg = DensePoseConfidenceModelConfig.from_cfg(cfg)
        self._initialize_confidence_estimation_layers(cfg, input_channels)
        self._registry = {}
        initialize_module_params(self)  # pyre-ignore[6]

    def _initialize_confidence_estimation_layers(self, cfg: CfgNode, dim_in: int):
        """
        Initialize confidence estimation layers based on configuration options

        Args:
            cfg (CfgNode): configuration options
            dim_in (int): number of input channels
        """
        kernel_size = cfg.MODEL.ROI_DENSEPOSE_HEAD.DECONV_KERNEL
        if self.confidence_model_cfg.segm_confidence.enabled:
            self.coarse_segm_confidence_lowres = ConvTranspose2d(  # pyre-ignore[16]
                dim_in, 1, kernel_size, stride=2, padding=int(kernel_size / 2 - 1)
            )

    def forward(self, head_outputs: torch.Tensor):
        """
        Perform forward operation on head outputs used as inputs for the predictor.
        Calls forward method from the base predictor and uses its outputs to compute
        confidences.

        Args:
            head_outputs (Tensor): head outputs used as predictor inputs
        Return:
            An instance of outputs with confidences,
            see `decorate_cse_predictor_output_class_with_confidences`
        """
        # assuming base class returns SIUV estimates in its first result
        base_predictor_outputs = super().forward(head_outputs)  # pyre-ignore[16]

        # create output instance by extending base predictor outputs:
        output = self._create_output_instance(base_predictor_outputs)

        if self.confidence_model_cfg.segm_confidence.enabled:
            # base predictor outputs are assumed to have `coarse_segm` attribute
            # base predictor is assumed to define `interp2d` method for bilinear interpolation
            output.coarse_segm_confidence = (
                F.softplus(
                    self.interp2d(  # pyre-ignore[16]
                        self.coarse_segm_confidence_lowres(head_outputs)  # pyre-ignore[16]
                    )
                )
                + self.confidence_model_cfg.segm_confidence.epsilon
            )
            output.coarse_segm = base_predictor_outputs.coarse_segm * torch.repeat_interleave(
                output.coarse_segm_confidence, base_predictor_outputs.coarse_segm.shape[1], dim=1
            )

        return output

    def _create_output_instance(self, base_predictor_outputs: Any):
        """
        Create an instance of predictor outputs by copying the outputs from the
        base predictor and initializing confidence

        Args:
            base_predictor_outputs: an instance of base predictor outputs
                (the outputs type is assumed to be a dataclass)
        Return:
           An instance of outputs with confidences
        """
        PredictorOutput = decorate_cse_predictor_output_class_with_confidences(
            type(base_predictor_outputs)  # pyre-ignore[6]
        )
        # base_predictor_outputs is assumed to be a dataclass
        # reassign all the fields from base_predictor_outputs (no deep copy!), add new fields
        output = PredictorOutput(
            **base_predictor_outputs.__dict__,
            coarse_segm_confidence=None,
        )
        return output