import torch


def obj_sampler(num_obj_samples,
                fg_mask,
                flatten_centerness_targets,
                flatten_gt_inds,
                *args,
                uniform_mix_ratio=0.5,
                eps=1e-5):
    """
    Sample from the dense objects to reduce the Monte Carlo overhead in training the EPro-PnP

    Args:
        num_obj_samples (int): Target sample count. If fg_mask_count == 0, the actual count
            will be zero
        fg_mask (torch.Tensor): Shape (num_total_point, ), bool
        flatten_centerness_targets (torch.Tensor): Shape (num_total_point, )
        flatten_gt_inds (torch.Tensor): Shape (num_total_point, ), int64
        *args (Tuple[torch.Tensor]): Shape (num_total_point, *), optional

    Returns:
        Tuple[torch.Tensor]:
            sample_gt_inds: Shape (num_obj_sample, ), int64
            sample_weights: Shape (num_obj_sample, ), balanced sample weights, where the
                cummulative weights of samples assiciated with each g.t. are equal
            sample_uniform_weights: Shape (num_obj_sample, ), balanced uniform sample weights,
                where the cummulative weights of samples assiciated with each g.t. are equal,
                and the weights of samples of the same g.t. are uniform. If uniform_mix_ratio
                == 0.0, then sample_uniform_weights and sample_weights are equivalent
            optional_outputs: Shape (num_obj_sample, *)
    """
    fg_mask_count = fg_mask.count_nonzero()

    num_uniform_samples = min(int(round(num_obj_samples * uniform_mix_ratio)), fg_mask_count)
    num_replace_samples = num_obj_samples - num_uniform_samples
    uniform_mix_ratio = num_uniform_samples / num_obj_samples

    # (num_total_point, )
    prob = flatten_centerness_targets * fg_mask
    prob /= prob.sum().clamp(min=eps)
    prob_uniform = fg_mask / fg_mask_count.clamp(min=1)  # uniform
    prob_mix = prob_uniform * uniform_mix_ratio + prob * (1 - uniform_mix_ratio)
    
    if fg_mask_count > 0:
        sample_point_inds_uniform = torch.multinomial(
            prob_uniform, num_uniform_samples, replacement=False)
        sample_point_inds_replace = torch.multinomial(
            prob, num_replace_samples, replacement=True)
        sample_point_inds = torch.cat(
            (sample_point_inds_uniform, sample_point_inds_replace), dim=0)
        sample_gt_inds = flatten_gt_inds[sample_point_inds]
        sample_prob_weights = prob[sample_point_inds] / prob_mix[sample_point_inds]
        
        gt_mask = sample_gt_inds == torch.arange(
            sample_gt_inds.max() + 1, device=sample_gt_inds.device)[:, None]  # (num_gt, num_obj_sample)
        
        gt_prob_sum = (sample_prob_weights * gt_mask).sum(dim=1)
        gt_weights = 1 / gt_prob_sum.clamp(min=eps)
        sample_weights = sample_prob_weights * gt_weights[sample_gt_inds]
        sample_weights /= sample_weights.mean().clamp(min=eps)
       
        gt_counts = torch.count_nonzero(gt_mask, dim=1)
        gt_weights = 1 / gt_counts.clamp(min=1)
        sample_uniform_weights = gt_weights[sample_gt_inds]
        sample_uniform_weights /= sample_uniform_weights.mean().clamp(min=eps)
        
    else:
        sample_point_inds = sample_gt_inds = flatten_gt_inds[[]]
        sample_weights = sample_uniform_weights = prob[[]]

    ret = (sample_gt_inds, sample_weights, sample_uniform_weights)
    extra_ret = []
    for arg in args:
        extra_ret.append(arg[sample_point_inds])
    return ret + tuple(extra_ret)