"""
Code borrowed from SelfMask: https://github.com/NoelShin/selfmask
"""

import torch


class FMeasure:
    def __init__(
        self,
        default_thres: float = 0.5,
        beta_square: float = 0.3,
        n_bins: int = 255,
        eps: float = 1e-7,
    ):
        """
        :param default_thres: a hyperparameter for F-measure that is used to binarize a predicted mask. Default: 0.5
        :param beta_square: a hyperparameter for F-measure. Default: 0.3
        :param n_bins: the number of thresholds that will be tested for F-max. Default: 255
        :param eps: a small value for numerical stability
        """

        self.beta_square = beta_square
        self.default_thres = default_thres
        self.eps = eps
        self.n_bins = n_bins

    def _compute_precision_recall(
        self, binary_pred_mask: torch.Tensor, gt_mask: torch.Tensor
    ) -> torch.Tensor:
        """
        :param binary_pred_mask: (B x H x W) or (H x W)
        :param gt_mask: (B x H x W) or (H x W), should be the same with binary_pred_mask
        """
        tp = torch.logical_and(binary_pred_mask, gt_mask).sum(dim=(-1, -2))
        tp_fp = binary_pred_mask.sum(dim=(-1, -2))
        tp_fn = gt_mask.sum(dim=(-1, -2))

        prec = tp / (tp_fp + self.eps)
        recall = tp / (tp_fn + self.eps)
        return prec, recall

    def _compute_f_measure(
        self,
        pred_mask: torch.Tensor,
        gt_mask: torch.Tensor,
        thresholds: torch.Tensor = None,
    ) -> torch.Tensor:
        if thresholds is None:
            binary_pred_mask = pred_mask > self.default_thres
        else:
            binary_pred_mask = pred_mask > thresholds

        prec, recall = self._compute_precision_recall(binary_pred_mask, gt_mask)
        f_measure = ((1 + (self.beta_square**2)) * prec * recall) / (
            (self.beta_square**2) * prec + recall + self.eps
        )
        return f_measure.cpu()

    def _compute_f_max(
        self, pred_mask: torch.Tensor, gt_mask: torch.Tensor
    ) -> torch.Tensor:
        """Compute self.n_bins + 1  F-measures, each of which has a different threshold, then return the maximum
        F-measure among them.

        :param pred_mask: (H x W)
        :param gt_mask: (H x W)
        """

        # pred_masks, gt_masks: H x W -> self.n_bins x H x W
        pred_masks = pred_mask.unsqueeze(dim=0).repeat(self.n_bins, 1, 1)
        gt_masks = gt_mask.unsqueeze(dim=0).repeat(self.n_bins, 1, 1)

        # thresholds: self.n_bins x 1 x 1
        thresholds = (
            torch.arange(0, 1, 1 / self.n_bins)
            .view(self.n_bins, 1, 1)
            .to(pred_masks.device)
        )

        # f_measures: self.n_bins
        f_measures = self._compute_f_measure(pred_masks, gt_masks, thresholds)
        return torch.max(f_measures).cpu(), f_measures

    def _compute_f_mean(
        self,
        pred_mask: torch.Tensor,
        gt_mask: torch.Tensor,
    ) -> torch.Tensor:
        adaptive_thres = 2 * pred_mask.mean(dim=(-1, -2), keepdim=True)
        binary_pred_mask = pred_mask > adaptive_thres

        prec, recall = self._compute_precision_recall(binary_pred_mask, gt_mask)
        f_mean = ((1 + (self.beta_square**2)) * prec * recall) / (
            (self.beta_square**2) * prec + recall + self.eps
        )
        return f_mean.cpu()

    def __call__(self, pred_mask: torch.Tensor, gt_mask: torch.Tensor) -> dict:
        """
        :param pred_mask: (H x W) a normalized prediction mask with values in [0, 1]
        :param gt_mask: (H x W) a binary ground truth mask with values in {0, 1}
        :return: a dictionary with keys being "f_measure" and "f_max" and values being the respective values.
        """
        outputs: dict = dict()
        for k in ("f_measure", "f_mean"):
            outputs.update({k: getattr(self, f"_compute_{k}")(pred_mask, gt_mask)})

        f_max_, all_f = self._compute_f_max(pred_mask, gt_mask)
        outputs["f_max"] = f_max_
        outputs["all_f"] = all_f  # List of all f values for all thresholds
        return outputs