import torch

def safe_log(z):
    return torch.log(z + 1e-7)

def log_sum_exp(value, dim=None, keepdim=False):
    """Numerically stable implementation of the operation
    value.exp().sum(dim, keepdim).log()
    """
    if dim is not None:
        m, _ = torch.max(value, dim=dim, keepdim=True)
        value0 = value - m
        if keepdim is False:
            m = m.squeeze(dim)
        return m + torch.log(torch.sum(torch.exp(value0), dim=dim, keepdim=keepdim))
    else:
        m = torch.max(value)
        sum_exp = torch.sum(torch.exp(value - m))
        return m + torch.log(sum_exp)

def exp_mean_log(x, factor=None, dim=1):
    # e.g. x = torch.Tensor([[-800, -900, -1000]
    #                        [-100, -200, -300 ]]) # [batch_size, nsamples]
    # value.exp().mean(dim).log(), when factor is None
    # (value.exp()*factor).mean(dim).log(), otherwise
    base, _ = torch.max(x, dim=dim)
    exp_offset = (x-base.unsqueeze(dim=dim)).exp()
    if factor is None:
        return exp_offset.mean(dim=dim).log() + base
    else:
        return (exp_offset*factor).mean(dim=dim).log() + base


def generate_grid(zmin, zmax, dz, device, ndim=2):
    """generate a 1- or 2-dimensional grid
    Returns: Tensor, int
        Tensor: The grid tensor with shape (k^2, 2),
            where k=(zmax - zmin)/dz
        int: k
    """

    if ndim == 2:
        x = torch.arange(zmin, zmax, dz)
        k = x.size(0)

        x1 = x.unsqueeze(1).repeat(1, k).view(-1)
        x2 = x.repeat(k)

        return torch.cat((x1.unsqueeze(-1), x2.unsqueeze(-1)), dim=-1).to(device), k

    elif ndim == 1:
        return torch.arange(zmin, zmax, dz).unsqueeze(1).to(device)