# 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 ...structures import decorate_predictor_output_class_with_confidences
from ..confidence import DensePoseConfidenceModelConfig, DensePoseUVConfidenceType
from ..utils import initialize_module_params


class DensePoseChartConfidencePredictorMixin:
    """
    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 segmentation and UV tensors estimated by some
    base predictor. Several assumptions need to hold for the base predictor:
    1) the `forward` method must return SIUV tuple as the first result (
        S = coarse segmentation, I = fine segmentation, U and V are intrinsic
        chart coordinates)
    2) `interp2d` method must be defined to perform bilinear interpolation;
        the same method is typically used for SIUV 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
        """
        dim_out_patches = cfg.MODEL.ROI_DENSEPOSE_HEAD.NUM_PATCHES + 1
        kernel_size = cfg.MODEL.ROI_DENSEPOSE_HEAD.DECONV_KERNEL
        if self.confidence_model_cfg.uv_confidence.enabled:
            if self.confidence_model_cfg.uv_confidence.type == DensePoseUVConfidenceType.IID_ISO:
                self.sigma_2_lowres = ConvTranspose2d(  # pyre-ignore[16]
                    dim_in, dim_out_patches, kernel_size, stride=2, padding=int(kernel_size / 2 - 1)
                )
            elif (
                self.confidence_model_cfg.uv_confidence.type
                == DensePoseUVConfidenceType.INDEP_ANISO
            ):
                self.sigma_2_lowres = ConvTranspose2d(
                    dim_in, dim_out_patches, kernel_size, stride=2, padding=int(kernel_size / 2 - 1)
                )
                self.kappa_u_lowres = ConvTranspose2d(  # pyre-ignore[16]
                    dim_in, dim_out_patches, kernel_size, stride=2, padding=int(kernel_size / 2 - 1)
                )
                self.kappa_v_lowres = ConvTranspose2d(  # pyre-ignore[16]
                    dim_in, dim_out_patches, kernel_size, stride=2, padding=int(kernel_size / 2 - 1)
                )
            else:
                raise ValueError(
                    f"Unknown confidence model type: "
                    f"{self.confidence_model_cfg.confidence_model_type}"
                )
        if self.confidence_model_cfg.segm_confidence.enabled:
            self.fine_segm_confidence_lowres = ConvTranspose2d(  # pyre-ignore[16]
                dim_in, 1, kernel_size, stride=2, padding=int(kernel_size / 2 - 1)
            )
            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_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.uv_confidence.enabled:
            if self.confidence_model_cfg.uv_confidence.type == DensePoseUVConfidenceType.IID_ISO:
                # assuming base class defines interp2d method for bilinear interpolation
                output.sigma_2 = self.interp2d(self.sigma_2_lowres(head_outputs))  # pyre-ignore[16]
            elif (
                self.confidence_model_cfg.uv_confidence.type
                == DensePoseUVConfidenceType.INDEP_ANISO
            ):
                # assuming base class defines interp2d method for bilinear interpolation
                output.sigma_2 = self.interp2d(self.sigma_2_lowres(head_outputs))
                output.kappa_u = self.interp2d(self.kappa_u_lowres(head_outputs))  # pyre-ignore[16]
                output.kappa_v = self.interp2d(self.kappa_v_lowres(head_outputs))  # pyre-ignore[16]
            else:
                raise ValueError(
                    f"Unknown confidence model type: "
                    f"{self.confidence_model_cfg.confidence_model_type}"
                )
        if self.confidence_model_cfg.segm_confidence.enabled:
            # base predictor outputs are assumed to have `fine_segm` and `coarse_segm` attributes
            # base predictor is assumed to define `interp2d` method for bilinear interpolation
            output.fine_segm_confidence = (
                F.softplus(
                    self.interp2d(self.fine_segm_confidence_lowres(head_outputs))  # pyre-ignore[16]
                )
                + self.confidence_model_cfg.segm_confidence.epsilon
            )
            output.fine_segm = base_predictor_outputs.fine_segm * torch.repeat_interleave(
                output.fine_segm_confidence, base_predictor_outputs.fine_segm.shape[1], dim=1
            )
            output.coarse_segm_confidence = (
                F.softplus(
                    self.interp2d(
                        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_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,
            fine_segm_confidence=None,
            sigma_1=None,
            sigma_2=None,
            kappa_u=None,
            kappa_v=None,
        )
        return output