import torch


def sample_pdf(bins, weights, Nsamples=64, det=False):
    """
    Sample @Nsamples depths from @bins with distribution defined by @weights.
    
    Inputs:
        bins: depth bins boundaries, (*Nrays, Nbins+1) where Nbins == Nsamples(coarse)-2
                              1          2           3            4     Nsamples(coarse)=5
            z_vals (coarse) : x          x           x            x           x
            bins (midpoints):      [           ][          ][           ]
                                         1           2         Nbins=3
        weights: (*Nrays, Nbins)        w_1         w_2        w_Nbins    
        Nsamples: the number of samples to draw from the distribution
        det: sampling deterministic or not
    Outputs:
        z_vals: the sampled depth values from bins, (*Nrays, Nsamples)

    NOTE: @weights should be detached before feeding into the function.
    """
    eps = 1e-5
    *Nrays, Nbins = weights.shape
    weights = weights + eps  # prevent division by zero (don't do inplace op!)
    pdf = weights / weights.sum(-1, keepdim=True)               # (*Nrays, Nbins)
    cdf = torch.cumsum(pdf, -1)                                 # (*Nrays, Nbins)
    cdf = torch.cat([torch.zeros_like(cdf[..., :1]), cdf], -1)  # (*Nrays, Nbins+1)
                                                                # Prepend 0 to CDF
    if det:
        u = torch.linspace(0, 1, Nsamples, device=bins.device)
        u = u.expand([*Nrays, Nsamples])
    else:
        u = torch.rand([*Nrays, Nsamples], device=bins.device)
    u = u.contiguous()

    # Find correct intervals for u
    inds = torch.searchsorted(cdf, u, right=True)  # (*Nrays, Nsamples)
    below = torch.clamp_min(inds-1, 0)             # (*Nrays, Nsamples)
    above = torch.clamp_max(inds, Nbins)           # (*Nrays, Nsamples)
    # print('below.shape: ', below.shape)
    # print('above.shape: ', above.shape)

    cdf_below = torch.gather(cdf, -1, below)       # (*Nrays, Nsamples)
    cdf_above = torch.gather(cdf, -1, above)       # (*Nrays, Nsamples)
    # print('cdf_below.shape: ', cdf_below.shape)
    # print('cdf_above.shape: ', cdf_above.shape)

    bins_below = torch.gather(bins, -1, below)     # (*Nrays, Nsamples)
    bins_above = torch.gather(bins, -1, above)     # (*Nrays, Nsamples)
    # print('bins_below.shape: ', bins_below.shape)
    # print('bins_above.shape: ', bins_above.shape)

    cdf_diff = cdf_above - cdf_below
    # print('cdf_diff.shape: ', cdf_diff.shape)
    cdf_diff[cdf_diff<eps] = 1.0

    t = (u - cdf_below) / cdf_diff
    z_vals = bins_below + t * (bins_above - bins_below)  # (*Nrays, Nsamples)
    return z_vals