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

import random
from typing import Optional, Tuple
import torch
from torch.nn import functional as F

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

from densepose.converters.base import IntTupleBox

from .densepose_cse_base import DensePoseCSEBaseSampler


class DensePoseCSEConfidenceBasedSampler(DensePoseCSEBaseSampler):
    """
    Samples DensePose data from DensePose predictions.
    Samples for each class are drawn using confidence value estimates.
    """

    def __init__(
        self,
        cfg: CfgNode,
        use_gt_categories: bool,
        embedder: torch.nn.Module,
        confidence_channel: str,
        count_per_class: int = 8,
        search_count_multiplier: Optional[float] = None,
        search_proportion: Optional[float] = None,
    ):
        """
        Constructor

        Args:
          cfg (CfgNode): the config of the model
          embedder (torch.nn.Module): necessary to compute mesh vertex embeddings
          confidence_channel (str): confidence channel to use for sampling;
            possible values:
              "coarse_segm_confidence": confidences for coarse segmentation
            (default: "coarse_segm_confidence")
          count_per_class (int): the sampler produces at most `count_per_class`
              samples for each category (default: 8)
          search_count_multiplier (float or None): if not None, the total number
              of the most confident estimates of a given class to consider is
              defined as `min(search_count_multiplier * count_per_class, N)`,
              where `N` is the total number of estimates of the class; cannot be
              specified together with `search_proportion` (default: None)
          search_proportion (float or None): if not None, the total number of the
              of the most confident estimates of a given class to consider is
              defined as `min(max(search_proportion * N, count_per_class), N)`,
              where `N` is the total number of estimates of the class; cannot be
              specified together with `search_count_multiplier` (default: None)
        """
        super().__init__(cfg, use_gt_categories, embedder, count_per_class)
        self.confidence_channel = confidence_channel
        self.search_count_multiplier = search_count_multiplier
        self.search_proportion = search_proportion
        assert (search_count_multiplier is None) or (search_proportion is None), (
            f"Cannot specify both search_count_multiplier (={search_count_multiplier})"
            f"and search_proportion (={search_proportion})"
        )

    def _produce_index_sample(self, values: torch.Tensor, count: int):
        """
        Produce a sample of indices to select data based on confidences

        Args:
            values (torch.Tensor): a tensor of length k that contains confidences
                k: number of points labeled with part_id
            count (int): number of samples to produce, should be positive and <= k

        Return:
            list(int): indices of values (along axis 1) selected as a sample
        """
        k = values.shape[1]
        if k == count:
            index_sample = list(range(k))
        else:
            # take the best count * search_count_multiplier pixels,
            # sample from them uniformly
            # (here best = smallest variance)
            _, sorted_confidence_indices = torch.sort(values[0])
            if self.search_count_multiplier is not None:
                search_count = min(int(count * self.search_count_multiplier), k)
            elif self.search_proportion is not None:
                search_count = min(max(int(k * self.search_proportion), count), k)
            else:
                search_count = min(count, k)
            sample_from_top = random.sample(range(search_count), count)
            index_sample = sorted_confidence_indices[-search_count:][sample_from_top]
        return index_sample

    def _produce_mask_and_results(
        self, instance: Instances, bbox_xywh: IntTupleBox
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Method to get labels and DensePose results from an instance

        Args:
            instance (Instances): an instance of
                `DensePoseEmbeddingPredictorOutputWithConfidences`
            bbox_xywh (IntTupleBox): the corresponding bounding box

        Return:
            mask (torch.Tensor): shape [H, W], DensePose segmentation mask
            embeddings (Tuple[torch.Tensor]): a tensor of shape [D, H, W]
                DensePose CSE Embeddings
            other_values: a tensor of shape [1, H, W], DensePose CSE confidence
        """
        _, _, w, h = bbox_xywh
        densepose_output = instance.pred_densepose
        mask, embeddings, _ = super()._produce_mask_and_results(instance, bbox_xywh)
        other_values = F.interpolate(
            getattr(densepose_output, self.confidence_channel),
            size=(h, w),
            mode="bilinear",
        )[0].cpu()
        return mask, embeddings, other_values