import torch


def argSoftmax3d(tensor, temperature, num_samples=1):
    """
  Find the index of the Softmax value in a 3d tensor. Prob is proportional to q/temperature

  Args:
    - tensor: PyTorch tensor of size (n x c x d x d)

  Returns:  n x num_samples x 2 PyTorch tensor containing indexes of max values at X Y
            n x num_samples x 1 PyTorch tensor containing indexes of max values at theta
  """
    n = tensor.size(0)
    c = tensor.size(1)
    d = tensor.size(2)
    probs = (tensor / temperature).view(n, -1).softmax(dim=-1)
    m = torch.multinomial(probs, num_samples)
    theta = m // (d * d)
    xy_idex = m - theta * d * d

    return torch.cat(((xy_idex // d).view(-1, 1), (xy_idex % d).view(-1, 1)), dim=1).long(), theta.view(-1, 1).long()