import torch


def get_pmf(prelim_assignment):
    """
    Calculates the probabilitiy mass function (pmf) which here is simplified the number of weights assigned to a
    specific centroid devided by th number of all weights. The preliminary assignment considers only the minimal
    distance from  centroids to weights as cost.
    With "spars_bound" we ensure that at least 50% of all weights would be assigned to the zero-centroid w_0 such that
    the entropy score (information content) for w_0 is always the lowest.

    Parameters:
    -----------
        prelim_assignment:
            Minimal arguments for all 3 centroid distances to all layer weights

    Returns:
    --------
        pmf_prelim:
            Percentage frequencies of -only distance dependent- centroid assignments (w_n, w_0, w_p)
            with pmf[w_n] + pmf[w_0] + pmf[w_p] = 1 and  pmf[w_0] always > 0.5
    """
    C_counts = torch.ones(3)
    C_val, C_cts = torch.unique(prelim_assignment, return_counts=True)

    # For the usual case that layer weights are assigned to three centroids (ternary)
    if C_cts.shape[0] == 3:
        C_counts = C_cts

    # The following two cases, especially the last one, should not occur as precautions were taken such that the
    # w_0 centroid can't absorb all assignments
    # If layer weights are assigned to two only centroids (binary)
    elif C_cts.shape[0] == 2:
        if 0 not in C_val:
            C_counts[1] = C_cts[0]
            C_counts[2] = C_cts[1]
        if 1 not in C_val:
            C_counts[0] = C_cts[0]
            C_counts[2] = C_cts[1]
        if 2 not in C_val:
            C_counts[0] = C_cts[0]
            C_counts[1] = C_cts[1]
    # If layer weights are assigned to only one centroid
    elif C_cts.shape[0] == 1:
        if (0 not in C_val and 1 not in C_val):
            C_counts[2] = C_cts[0]
        if (0 not in C_val and 2 not in C_val):
            C_counts[1] = C_cts[0]
        if (1 not in C_val and 2 not in C_val):
            C_counts[0] = C_cts[0]

    pmf_prelim = torch.div(C_counts.type(torch.float32), torch.numel(prelim_assignment))

    # Ensuring that at least 50% of all weights are assigned to w_0 and probabilities still sum up to 1
    spars_bound = 0.5
    if pmf_prelim[1] < spars_bound:
        pmf_prelim[0] -= (pmf_prelim[0]/(pmf_prelim[0] +
                                        pmf_prelim[2]))*(spars_bound - pmf_prelim[1])
        pmf_prelim[2] -= (pmf_prelim[2] / (pmf_prelim[0] +
                                           pmf_prelim[2])) * (spars_bound - pmf_prelim[1])
        pmf_prelim[1] = spars_bound

    return pmf_prelim