from typing import Union
import torch
from torch.distributions import Categorical, Independent, Normal


def compute_importance_weights(
    target_output: Union[torch.Tensor, dict],
    behaviour_output: Union[torch.Tensor, dict],
    action: torch.Tensor,
    action_space_type: str = 'discrete',
    requires_grad: bool = False
):
    """
    Overview:
        Computing importance sampling weight with given output and action
    Arguments:
        - target_output (:obj:`Union[torch.Tensor,dict]`): the output taking the action \
            by the current policy network, \
            usually this output is network output logit if action space is discrete, \
            or is a dict containing parameters of action distribution if action space is continuous.
        - behaviour_output (:obj:`Union[torch.Tensor,dict]`): the output taking the action \
            by the behaviour policy network,\
            usually this output is network output logit,  if action space is discrete, \
            or is a dict containing parameters of action distribution if action space is continuous.
        - action (:obj:`torch.Tensor`): the chosen action(index for the discrete action space) in trajectory,\
            i.e.: behaviour_action
        - action_space_type (:obj:`str`): action space types in ['discrete', 'continuous']
        - requires_grad (:obj:`bool`): whether requires grad computation
    Returns:
        - rhos (:obj:`torch.Tensor`): Importance sampling weight
    Shapes:
        - target_output (:obj:`Union[torch.FloatTensor,dict]`): :math:`(T, B, N)`, \
            where T is timestep, B is batch size and N is action dim
        - behaviour_output (:obj:`Union[torch.FloatTensor,dict]`): :math:`(T, B, N)`
        - action (:obj:`torch.LongTensor`): :math:`(T, B)`
        - rhos (:obj:`torch.FloatTensor`): :math:`(T, B)`
    Examples:
        >>> target_output = torch.randn(2, 3, 4)
        >>> behaviour_output = torch.randn(2, 3, 4)
        >>> action = torch.randint(0, 4, (2, 3))
        >>> rhos = compute_importance_weights(target_output, behaviour_output, action)
    """
    grad_context = torch.enable_grad() if requires_grad else torch.no_grad()
    assert isinstance(action, torch.Tensor)
    assert action_space_type in ['discrete', 'continuous']

    with grad_context:
        if action_space_type == 'continuous':
            dist_target = Independent(Normal(loc=target_output['mu'], scale=target_output['sigma']), 1)
            dist_behaviour = Independent(Normal(loc=behaviour_output['mu'], scale=behaviour_output['sigma']), 1)
            rhos = dist_target.log_prob(action) - dist_behaviour.log_prob(action)
            rhos = torch.exp(rhos)
            return rhos
        elif action_space_type == 'discrete':
            dist_target = Categorical(logits=target_output)
            dist_behaviour = Categorical(logits=behaviour_output)
            rhos = dist_target.log_prob(action) - dist_behaviour.log_prob(action)
            rhos = torch.exp(rhos)
            return rhos