import numpy as np 

def get_uniform_mass_bins(probs, n_bins):
    assert(probs.size >= n_bins), "Fewer points than bins"
    
    probs_sorted = np.sort(probs)

    # split probabilities into groups of approx equal size
    groups = np.array_split(probs_sorted, n_bins)
    bin_edges = list()
    bin_upper_edges = list()

    for cur_group in range(n_bins-1):
        bin_upper_edges += [max(groups[cur_group])]
    bin_upper_edges += [np.inf]

    return np.array(bin_upper_edges)


def bin_points(scores, bin_edges):
    assert(bin_edges is not None), "Bins have not been defined"
    assert(scores.shape[0] == scores.size), "scores should be a 1D vector"
    scores = np.reshape(scores, (scores.size, 1))
    bin_edges = np.reshape(bin_edges, (1, bin_edges.size))
    return np.sum(scores > bin_edges, axis=1)

def bin_points_uniform(x, n_bins):
    x = x.squeeze()
    n_points = x.size
    assert(n_points >= n_bins), "Fewer points than bins"

    groups = np.array_split(probs_sorted, n_bins)
    bin_upper_edges = list()
    for cur_group in range(n_bins-1):
        bin_edges += [(max(groups[cur_group])+min(groups[cur_group+1]))/2]
        bin_upper_edges += [max(groups[cur_group])]
    bin_upper_edges += [np.inf]

    '''
    bin_upper_edges = np.interp(np.linspace(0, n_points, n_bins + 1),
                            np.arange(n_points),
                                np.sort(x))[1:]
    bin_upper_edges[n_bins - 1] = np.inf
    '''
    return np.sum(x.reshape((-1, 1)) > bin_upper_edges, axis=1)
