utils

class Interpolate(nn.Module):
    """Wrapper for torch.nn.functional.interpolate."""

    def __init__(self, size=None, scale=None, mode="bilinear", align_corners=False):
        super().__init__()
        assert (size is None) == (scale is not None)
        self.size = size
        self.scale = scale
        self.mode = mode
        self.align_corners = align_corners

    def forward(self, x):
        out = F.interpolate(
            x,
            size=self.size,
            scale_factor=self.scale,
            mode=self.mode,
            align_corners=self.align_corners,
        )
        return out

class StableExponential:
    """
    Class that redefines the definition of exp() to increase numerical stability.
    Naturally, also the definition of log() must change accordingly.
    However, it is worth noting that the two operations remain one the inverse of the other,
    meaning that x = log(exp(x)) and x = exp(log(x)) are always true.

    Definition:
        exp(x) = {
            exp(x) if x<=0
            x+1    if x>0
        }

        log(x) = {
            x        if x<=0
            log(1+x) if x>0
        }

    NOTE 1:
        Within the class everything is done on the tensor given as input to the constructor.
        Therefore, when exp() is called, self._tensor.exp() is computed.
        When log() is called, torch.log(self._tensor.exp()) is computed instead.

    NOTE 2:
        Given the output from exp(), torch.log() or the log() method of the class give identical results.
    """

    def __init__(self, tensor):
        self._raw_tensor = tensor
        posneg_dic = self.posneg_separation(self._raw_tensor)
        self.pos_f, self.neg_f = posneg_dic["filter"]
        self.pos_data, self.neg_data = posneg_dic["value"]

    def posneg_separation(self, tensor):
        pos = tensor > 0
        pos_tensor = torch.clip(tensor, min=0)

        neg = tensor <= 0
        neg_tensor = torch.clip(tensor, max=0)

        return {"filter": [pos, neg], "value": [pos_tensor, neg_tensor]}

    def exp(self):
        return torch.exp(self.neg_data) * self.neg_f + (1 + self.pos_data) * self.pos_f

    def log(self):
        return self.neg_data * self.neg_f + torch.log(1 + self.pos_data) * self.pos_f

class StableLogVar:
    """
    Class that provides a numerically stable implementation of Log-Variance.
    Specifically, it uses the exp() and log() formulas defined in `StableExponential` class.
    """

    def __init__(
        self, logvar: torch.Tensor, enable_stable: bool = True, var_eps: float = 1e-6
    ):
        """
        Constructor.

        Parameters
        ----------
        logvar: torch.Tensor
            The input (true) logvar vector, to be converted in the Stable version.
        enable_stable: bool, optional
            Whether to compute the stable version of log-variance. Default is `True`.
        var_eps: float, optional
            The minimum value attainable by the variance. Default is `1e-6`.
        """
        self._lv = logvar
        self._enable_stable = enable_stable
        self._eps = var_eps

    def get(self) -> torch.Tensor:
        if self._enable_stable is False:
            return self._lv

        return torch.log(self.get_var())

    def get_var(self) -> torch.Tensor:
        """
        Get Variance from Log-Variance.
        """
        if self._enable_stable is False:
            return torch.exp(self._lv)
        return StableExponential(self._lv).exp() + self._eps

    def get_std(self) -> torch.Tensor:
        return torch.sqrt(self.get_var())

    @property
    def is_3D(self) -> bool:
        """Check if the _lv tensor is 3D.

        Recall that, in this framework, tensors have shape (B, C, [Z], Y, X).
        """
        return self._lv.dim() == 5

    def centercrop_to_size(self, size: Sequence[int]) -> None:
        """
        Centercrop the log-variance tensor to the desired size.

        Parameters
        ----------
        size: torch.Tensor
            The desired size of the log-variance tensor.
        """
        assert not self.is_3D, "Centercrop is implemented only for 2D tensors."

        if self._lv.shape[-1] == size:
            return

        diff = self._lv.shape[-1] - size
        assert diff > 0 and diff % 2 == 0
        self._lv = F.center_crop(self._lv, (size, size))

class StableMean:

    def __init__(self, mean):
        self._mean = mean

    def get(self) -> torch.Tensor:
        return self._mean

    @property
    def is_3D(self) -> bool:
        """Check if the _mean tensor is 3D.

        Recall that, in this framework, tensors have shape (B, C, [Z], Y, X).
        """
        return self._mean.dim() == 5

    def centercrop_to_size(self, size: Sequence[int]) -> None:
        """Centercrop the mean tensor to the desired size.

        Implemented only in the case of 2D tensors.

        Parameters
        ----------
        size: torch.Tensor
            The desired size of the log-variance tensor.
        """
        assert not self.is_3D, "Centercrop is implemented only for 2D tensors."

        if self._mean.shape[-1] == size:
            return

        diff = self._mean.shape[-1] - size
        assert diff > 0 and diff % 2 == 0
        self._mean = F.center_crop(self._mean, (size, size))