import torch
import math


def calculate_log_pi(log_stds: torch.Tensor, noises: torch.Tensor, actions: torch.Tensor) -> torch.Tensor:
    """
    Calculate log(\pi(a|s)) given log(std) of the distribution, noises, and actions to take

    Parameters
    ----------
    log_stds: torch.Tensor
        log(std) of the distribution
    noises: torch.Tensor
        noises added to the action
    actions: torch.Tensor
        actions to take

    Returns
    -------
    log_pi: torch.Tensor
        log(\pi(a|s))
    """
    gaussian_log_probs = (-0.5 * noises.pow(2) - log_stds).sum(
        dim=-1, keepdim=True) - 0.5 * math.log(2 * math.pi) * log_stds.size(-1)

    return gaussian_log_probs - torch.log(
        1 - actions.pow(2) + 1e-6).sum(dim=-1, keepdim=True)