likelihoods

class GaussianLikelihood(LikelihoodModule):
    r"""A specialized `LikelihoodModule` for Gaussian likelihood.

    Specifically, in the LVAE model, the likelihood is defined as:
        p(x|z_1) = N(x|\mu_{p,1}, \sigma_{p,1}^2)
    """

    def __init__(
        self,
        predict_logvar: Union[Literal["pixelwise"], None] = None,
        logvar_lowerbound: Union[float, None] = None,
    ):
        """Constructor.

        Parameters
        ----------
        predict_logvar: Union[Literal["pixelwise"], None], optional
            If `pixelwise`, log-variance is computed for each pixel, else log-variance
            is not computed. Default is `None`.
        logvar_lowerbound: float, optional
            The lowerbound value for log-variance. Default is `None`.
        """
        super().__init__()

        self.predict_logvar = predict_logvar
        self.logvar_lowerbound = logvar_lowerbound
        assert self.predict_logvar in [None, "pixelwise"]

        print(
            f"[{self.__class__.__name__}] PredLVar:{self.predict_logvar} LowBLVar:{self.logvar_lowerbound}"
        )

    def get_mean_lv(
        self, x: torch.Tensor
    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
        """
        Given the output of the top-down pass, compute the mean and log-variance of the
        Gaussian distribution defining the likelihood.

        Parameters
        ----------
        x: torch.Tensor
            The input tensor to the likelihood module, i.e., the output of the top-down
            pass.

        Returns
        -------
        tuple of (torch.tensor, optional torch.tensor)
            The first element of the tuple is the mean, the second element is the
            log-variance. If the attribute `predict_logvar` is `None` then the second
            element will be `None`.
        """

        # if LadderVAE.predict_logvar is None, dim 1 of `x`` has no. of target channels
        if self.predict_logvar is None:
            return x, None

        # Get pixel-wise mean and logvar
        # if LadderVAE.predict_logvar is not None,
        #   dim 1 has double no. of target channels
        mean, lv = x.chunk(2, dim=1)

        # Optionally, clip log-var to a lower bound
        if self.logvar_lowerbound is not None:
            lv = torch.clip(lv, min=self.logvar_lowerbound)

        return mean, lv

    def distr_params(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
        """
        Get parameters (mean, log-var) of the Gaussian distribution defined by the likelihood.

        Parameters
        ----------
        x: torch.Tensor
            The input tensor to the likelihood module, i.e., the output
            the LVAE 'output_layer'. Shape is: (B, 2 * C, [Z], Y, X) in case
            `predict_logvar` is not None, or (B, C, [Z], Y, X) otherwise.
        """
        mean, lv = self.get_mean_lv(x)
        params = {
            "mean": mean,
            "logvar": lv,
        }
        return params

    @staticmethod
    def mean(params: dict[str, torch.Tensor]) -> torch.Tensor:
        return params["mean"]

    @staticmethod
    def mode(params: dict[str, torch.Tensor]) -> torch.Tensor:
        return params["mean"]

    @staticmethod
    def sample(params: dict[str, torch.Tensor]) -> torch.Tensor:
        # p = Normal(params['mean'], (params['logvar'] / 2).exp())
        # return p.rsample()
        return params["mean"]

    @staticmethod
    def logvar(params: dict[str, torch.Tensor]) -> torch.Tensor:
        return params["logvar"]

    def log_likelihood(
        self, x: torch.Tensor, params: dict[str, Union[torch.Tensor, None]]
    ):
        """Compute Gaussian log-likelihood

        Parameters
        ----------
        x: torch.Tensor
            The target tensor. Shape is (B, C, [Z], Y, X).
        params: dict[str, Union[torch.Tensor, None]]
            The tensors obtained by chunking the output of the top-down pass,
            here used as parameters of the Gaussian distribution.

        Returns
        -------
        torch.Tensor
            The log-likelihood tensor. Shape is (B, C, [Z], Y, X).
        """
        if self.predict_logvar is not None:
            logprob = log_normal(x, params["mean"], params["logvar"])
        else:
            logprob = -0.5 * (params["mean"] - x) ** 2
        return logprob

class LikelihoodModule(nn.Module):
    """
    The base class for all likelihood modules.
    It defines the fundamental structure and methods for specialized likelihood models.
    """

    def distr_params(self, x: Any) -> None:
        return None

    def set_params_to_same_device_as(self, correct_device_tensor: Any) -> None:
        pass

    @staticmethod
    def logvar(params: Any) -> None:
        return None

    @staticmethod
    def mean(params: Any) -> None:
        return None

    @staticmethod
    def mode(params: Any) -> None:
        return None

    @staticmethod
    def sample(params: Any) -> None:
        return None

    def log_likelihood(self, x: Any, params: Any) -> None:
        return None

    def get_mean_lv(
        self, x: torch.Tensor
    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: ...

    def forward(
        self, input_: torch.Tensor, x: Union[torch.Tensor, None]
    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
        """
        Parameters:
        -----------
        input_: torch.Tensor
            The output of the top-down pass (e.g., reconstructed image in HDN,
            or the unmixed images in 'Split' models).
        x: Union[torch.Tensor, None]
            The target tensor. If None, the log-likelihood is not computed.
        """
        distr_params = self.distr_params(input_)
        mean = self.mean(distr_params)
        mode = self.mode(distr_params)
        sample = self.sample(distr_params)
        logvar = self.logvar(distr_params)

        if x is None:
            ll = None
        else:
            ll = self.log_likelihood(x, distr_params)

        dct = {
            "mean": mean,
            "mode": mode,
            "sample": sample,
            "params": distr_params,
            "logvar": logvar,
        }

        return ll, dct

class NoiseModelLikelihood(LikelihoodModule):

    def __init__(
        self,
        data_mean: Union[np.ndarray, torch.Tensor],
        data_std: Union[np.ndarray, torch.Tensor],
        noiseModel: NoiseModel,
    ):
        """Constructor.

        Parameters
        ----------
        data_mean: Union[np.ndarray, torch.Tensor]
            The mean of the data, used to unnormalize data for noise model evaluation.
        data_std: Union[np.ndarray, torch.Tensor]
            The standard deviation of the data, used to unnormalize data for noise
            model evaluation.
        noiseModel: NoiseModel
            The noise model instance used to compute the likelihood.
        """
        super().__init__()
        self.data_mean = torch.Tensor(data_mean)
        self.data_std = torch.Tensor(data_std)
        self.noiseModel = noiseModel

    def _set_params_to_same_device_as(
        self, correct_device_tensor: torch.Tensor
    ) -> None:
        """Set the parameters to the same device as the input tensor.

        Parameters
        ----------
        correct_device_tensor: torch.Tensor
            The tensor whose device is used to set the parameters.
        """
        if self.data_mean.device != correct_device_tensor.device:
            self.data_mean = self.data_mean.to(correct_device_tensor.device)
            self.data_std = self.data_std.to(correct_device_tensor.device)

    def get_mean_lv(self, x: torch.Tensor) -> tuple[torch.Tensor, None]:
        return x, None

    def distr_params(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
        mean, lv = self.get_mean_lv(x)
        params = {
            "mean": mean,
            "logvar": lv,
        }
        return params

    @staticmethod
    def mean(params: dict[str, torch.Tensor]) -> torch.Tensor:
        return params["mean"]

    @staticmethod
    def mode(params: dict[str, torch.Tensor]) -> torch.Tensor:
        return params["mean"]

    @staticmethod
    def sample(params: dict[str, torch.Tensor]) -> torch.Tensor:
        return params["mean"]

    def log_likelihood(self, x: torch.Tensor, params: dict[str, torch.Tensor]):
        """Compute the log-likelihood given the parameters `params` obtained
        from the reconstruction tensor and the target tensor `x`.

        Parameters
        ----------
        x: torch.Tensor
            The target tensor. Shape is (B, C, [Z], Y, X).
        params: dict[str, Union[torch.Tensor, None]]
            The tensors obtained from output of the top-down pass.
            Here, "mean" correspond to the whole output, while logvar is `None`.

        Returns
        -------
        torch.Tensor
            The log-likelihood tensor. Shape is (B, C, [Z], Y, X).
        """
        self._set_params_to_same_device_as(x)
        predicted_s_denormalized = params["mean"] * self.data_std + self.data_mean
        x_denormalized = x * self.data_std + self.data_mean
        likelihoods = self.noiseModel.likelihood(
            x_denormalized, predicted_s_denormalized
        )
        logprob = torch.log(likelihoods)
        return logprob