import torch


def downsample(vox_in, n, use_max=True):
    """
    Downsample a 3-d tensor n times
    Inputs:
      - vox_in (Tensor): HxWxD tensor
      - n (int): number of times to downsample each dimension
      - use_max (bool): use maximum value when downsampling. If set to False
                        the mean value is used.
    Output:
      - vox_out (Tensor): (H/n)x(W/n)x(D/n) tensor
    """
    dimy = vox_in.size(0) // n
    dimx = vox_in.size(1) // n
    dimz = vox_in.size(2) // n
    vox_out = torch.zeros((dimy, dimx, dimz))
    for x in range(dimx):
        for y in range(dimy):
            for z in range(dimz):
                subx = x * n
                suby = y * n
                subz = z * n
                subvox = vox_in[suby : suby + n, subx : subx + n, subz : subz + n]
                if use_max:
                    vox_out[y, x, z] = torch.max(subvox)
                else:
                    vox_out[y, x, z] = torch.mean(subvox)
    return vox_out