import torch


def sample_pdf(rays, weights, N_importance, det=False, eps=1e-5):
    """
    Sample @N_importance samples from @bins with distribution defined by @weights.
    from https://github.com/sxyu/pixel-nerf/blob/master/src/render/nerf.py#L120
    Inputs:
        rays: (N_rays, N_samples_+1) where N_samples_ is "the number of coarse samples per ray - 2"
        weights: (N_rays, N_samples_)
        N_importance: the number of samples to draw from the distribution
        det: deterministic or not
        eps: a small number to prevent division by zero

    Outputs:
        samples: the sampled samples
    """
    # torch.cuda.empty_cache()
    N_rays, N_samples_ = weights.shape
    weights = weights + eps # prevent division by zero (don't do inplace op!)
    pdf = weights / torch.sum(weights, -1, keepdim=True) # (N_rays, N_samples_)
    cdf = torch.cumsum(pdf, -1) # (N_rays, N_samples), cumulative distribution function
    cdf = torch.cat([torch.zeros_like(cdf[: ,:1]), cdf], -1)  # (N_rays, N_samples_+1) 
                                                               # padded to 0~1 inclusive
                                                               
    u = torch.rand(N_rays, N_importance, device=rays.device) # (N_rays, N_samples)
    # inds = searchsorted(cdf, u, side='right').float() - 1.0 # (N_rays, N_samples)
    inds = torch.searchsorted(cdf, u, right=True).float() - 1.0  # (B, Kf)
    inds = torch.clamp_min(inds, 0.0)

    z_steps = (inds + torch.rand_like(inds)) / N_samples_  # (B, Kf)

    near, far = rays[:, -2:-1], rays[:, -1:]  # (B, 1)
    # if not self.lindisp:  # Use linear sampling in depth space
    z_samp = near * (1 - z_steps) + far * z_steps  # (B, Kf)
    # else:  # Use linear sampling in disparity space
        # z_samp = 1 / (1 / near * (1 - z_steps) + 1 / far * z_steps)  # (B, Kf)
    return z_samp