import torch


def is_active(coords, active_voxels, volume_resolution):
    """ compute whether coordinates belong to active voxels by two criterias:
    1) a coordinate should be within the volume bound
    2) it belongs to the input active voxels
    
    Args:
        coords: [b, n_pts, n_steps, 3]
        active_voxels: [H, W, D]
    Returns:
        active_masks: [b, n_pts, n_steps]
    """

    batch, n_pts, n_steps = coords.shape[:3]
    within_grid = (coords[..., 0] >= 0) * (coords[..., 0] < volume_resolution[0]) *\
        (coords[..., 1] >= 0) * (coords[..., 1] < volume_resolution[1]) *\
        (coords[..., 2] >= 0) * (coords[..., 2] < volume_resolution[2])

    # 1 unit of voxel size for buffer
    active_x = (coords[..., 0] > -1) * (coords[..., 0] < volume_resolution[0] + 1)
    active_y = (coords[..., 1] > -1) * (coords[..., 1] < volume_resolution[1] + 1)
    active_z = (coords[..., 2] > -1) * (coords[..., 2] < volume_resolution[2] + 1)

    capped_x = torch.clip(coords[..., 0] * 1., 0, volume_resolution[0]-1).long().reshape(-1)
    capped_y = torch.clip(coords[..., 1] * 1., 0, volume_resolution[1]-1).long().reshape(-1)
    capped_z = torch.clip(coords[..., 2] * 1., 0, volume_resolution[2]-1).long().reshape(-1)

    active_masks = active_voxels[capped_x, capped_y, capped_z].reshape(batch, n_pts, n_steps)
    active_masks = torch.where(within_grid, active_masks, torch.ones_like(active_masks))
    active_masks = active_masks * active_x * active_y * active_z
    return active_masks