"""
This module contains generic VGAE functionalities, added as a Mixin to the
Variational Gene Program Graph Autoencoder neural network module.
"""

import torch


class VGAEModuleMixin:
    """
    VGAE module mix in class containing universal VGAE module 
    functionalities.
    """
    def reparameterize(self,
                       mu: torch.Tensor,
                       logstd: torch.Tensor) -> torch.Tensor:
        """
        Use reparameterization trick for latent space normal distribution.
        
        Parameters
        ----------
        mu:
            Expected values of the latent space distribution (dim: n_obs, 
            n_gps).
        logstd:
            Log standard deviations of the latent space distribution (dim: n_obs,
            n_gps).

        Returns
        ----------
        rep:
            Reparameterized latent features (dim: n_obs, n_gps).
        """
        if self.training:
            std = torch.exp(logstd)
            eps = torch.randn_like(mu)
            rep = eps.mul(std).add(mu)
            return rep
        else:
            rep = mu
            return rep