# Path: pytorchrl/distributions/base.py
# class Distribution(object):
# 
# 
#     def kl_sym(self, old_dist_info_vars, new_dist_info_vars):
#         """
#         Compute the symbolic KL divergence of two distributions
#         """
#         raise NotImplementedError
# 
#     def kl(self, old_dist_info, new_dist_info):
#         """
#         Compute the KL divergence of two distributions
#         """
#         raise NotImplementedError
# 
#     def likelihood_ratio(self, other, a):
#         """
#         Compute p_self(a) / p_other(a)
# 
#         Parameters
#         ----------
#         other (Distribution):
#         a (Variable):
# 
#         Returns
#         -------
#         ratio (Variable):
#         """
#         logli = self.log_likelihood(a)
#         other_logli = other.log_likelihood(a)
#         return (logli - other_logli).exp()
# 
#     def likelihood_ratio_sym(self, x_var, old_dist_info_vars, new_dist_info_vars):
#         raise NotImplementedError
# 
#     def entropy(self, dist_info):
#         raise NotImplementedError
# 
#     def log_likelihood_sym(self, x_var, dist_info_vars):
#         raise NotImplementedError
# 
#     def likelihood_sym(self, x_var, dist_info_vars):
#         raise NotImplementedError
# 
#     def log_likelihood(self, a):
#         raise NotImplementedError
# 
#     @property
#     def dist_info_keys(self):
#         raise NotImplementedError
# 
# Path: pytorchrl/misc/tensor_utils.py
# def constant(value):
#     """
#     Return a torch Variable for computation. This is a function to help
#     write short code.
# 
#     pytorch require multiplication take either two variables
#     or two tensors. And it is recommended to wrap a constant
#     in a variable which is kind of silly to me.
#     https://discuss.pytorch.org/t/adding-a-scalar/218
# 
#     Parameters
#     ----------
#     value (float): The value to be wrapped in Variable
# 
#     Returns
#     -------
#     constant (Variable): The Variable wrapped the value for computation.
#     """
#     return Variable(torch.Tensor([value])).type(torch.FloatTensor)


import numpy as np
import torch
from pytorchrl.distributions.base import Distribution
from pytorchrl.misc.tensor_utils import constant

        inverse of diagonal matrix is the diagonal matrix of elements at diagonal inverted, thus
        - (μ_2 - μ_1)^T Σ_2^{-1} (μ_2 - μ_1) = sum((means_1 - means_2)^2 / vars_2, axis=-1)

        trace is sum of the diagonal elements
        - tr(Σ_2^{-1}Σ_1) = sum(vars_1 / vars_2, axis=-1)

        Where

        - vars_1 = exp(2 * log_stds_1)

        - vars_2 = exp(2 * log_stds_2)

        Combined together, we have

        D_KL(N1||N2)
        = 1/2 ( tr(Σ_2^{-1}Σ_1) + (μ_2 - μ_1)^T Σ_2^{-1} (μ_2 - μ_1) - n + ln(det(Σ_2) / det(Σ_1)) )
        = sum(1/2 * ((vars_1 - vars_2) / vars_2 + (means_1 - means_2)^2 / vars_2 + 2 * (log_stds_2 - log_stds_1)), axis=-1)
        = sum( ((means_1 - means_2)^2 + vars_1 - vars_2) / (2 * vars_2) + (log_stds_2 - log_stds_1)), axis=-1)

        Parameters
        ----------
        other (DiagonalGaussian):

        Returns
        -------
        kl_div (Variable):
        """
        # Constant should wrap in Variable to multiply with another Variable
        # TODO (ewei) kl seems have problem

# Path: pytorchrl/distributions/base.py
# class Distribution(object):
# 
# 
#     def kl_sym(self, old_dist_info_vars, new_dist_info_vars):
#         """
#         Compute the symbolic KL divergence of two distributions
#         """
#         raise NotImplementedError
# 
#     def kl(self, old_dist_info, new_dist_info):
#         """
#         Compute the KL divergence of two distributions
#         """
#         raise NotImplementedError
# 
#     def likelihood_ratio(self, other, a):
#         """
#         Compute p_self(a) / p_other(a)
# 
#         Parameters
#         ----------
#         other (Distribution):
#         a (Variable):
# 
#         Returns
#         -------
#         ratio (Variable):
#         """
#         logli = self.log_likelihood(a)
#         other_logli = other.log_likelihood(a)
#         return (logli - other_logli).exp()
# 
#     def likelihood_ratio_sym(self, x_var, old_dist_info_vars, new_dist_info_vars):
#         raise NotImplementedError
# 
#     def entropy(self, dist_info):
#         raise NotImplementedError
# 
#     def log_likelihood_sym(self, x_var, dist_info_vars):
#         raise NotImplementedError
# 
#     def likelihood_sym(self, x_var, dist_info_vars):
#         raise NotImplementedError
# 
#     def log_likelihood(self, a):
#         raise NotImplementedError
# 
#     @property
#     def dist_info_keys(self):
#         raise NotImplementedError
# 
# Path: pytorchrl/misc/tensor_utils.py
# def constant(value):
#     """
#     Return a torch Variable for computation. This is a function to help
#     write short code.
# 
#     pytorch require multiplication take either two variables
#     or two tensors. And it is recommended to wrap a constant
#     in a variable which is kind of silly to me.
#     https://discuss.pytorch.org/t/adding-a-scalar/218
# 
#     Parameters
#     ----------
#     value (float): The value to be wrapped in Variable
# 
#     Returns
#     -------
#     constant (Variable): The Variable wrapped the value for computation.
#     """
#     return Variable(torch.Tensor([value])).type(torch.FloatTensor)


# Path: pytorchrl/distributions/diagonal_gaussian.py
import numpy as np
import torch
from pytorchrl.distributions.base import Distribution
from pytorchrl.misc.tensor_utils import constant


        inverse of diagonal matrix is the diagonal matrix of elements at diagonal inverted, thus
        - (μ_2 - μ_1)^T Σ_2^{-1} (μ_2 - μ_1) = sum((means_1 - means_2)^2 / vars_2, axis=-1)

        trace is sum of the diagonal elements
        - tr(Σ_2^{-1}Σ_1) = sum(vars_1 / vars_2, axis=-1)

        Where

        - vars_1 = exp(2 * log_stds_1)

        - vars_2 = exp(2 * log_stds_2)

        Combined together, we have

        D_KL(N1||N2)
        = 1/2 ( tr(Σ_2^{-1}Σ_1) + (μ_2 - μ_1)^T Σ_2^{-1} (μ_2 - μ_1) - n + ln(det(Σ_2) / det(Σ_1)) )
        = sum(1/2 * ((vars_1 - vars_2) / vars_2 + (means_1 - means_2)^2 / vars_2 + 2 * (log_stds_2 - log_stds_1)), axis=-1)
        = sum( ((means_1 - means_2)^2 + vars_1 - vars_2) / (2 * vars_2) + (log_stds_2 - log_stds_1)), axis=-1)

        Parameters
        ----------
        other (DiagonalGaussian):

        Returns
        -------
        kl_div (Variable):
        """
        # Constant should wrap in Variable to multiply with another Variable
        # TODO (ewei) kl seems have problem

        variance = (constant(2.0) * self.log_stds).exp()