import torch
import math


def sample_random_lightdirs(num_rays, num_samples, upper_only=False):
    """Randomly sample directions in the unit sphere.

    Args:
        num_rays: int or tensor shape dimension. Number of rays.
        num_samples: int or tensor shape dimension. Number of samples per ray.
        upper_only: bool. Whether to sample only on the upper hemisphere.

    Returns:
        lightdirs: [R, S, 3] float tensor. Random light directions sampled from the unit
            sphere for each sampled point.
    """
    if upper_only:
        min_z = 0
    else:
        min_z = -1

    phi = torch.rand(num_rays, num_samples) * (2 * math.pi)  # [R, S]
    cos_theta = torch.rand(num_rays, num_samples) * (1 - min_z) + min_z  # [R, S]
    theta = torch.acos(cos_theta)  # [R, S]

    x = torch.sin(theta) * torch.cos(phi)
    y = torch.sin(theta) * torch.sin(phi)
    z = torch.cos(theta)

    lightdirs = torch.cat((x[..., None], y[..., None], z[..., None]), dim=-1)  # [R, S, 3]
    return lightdirs