from typing import OrderedDict
import math


def calc_multinomial_sampling_prob_with_penalty(dataset_size, alpha=.5):
    """
    Calculate multinomial probability distribution based on https://arxiv.org/pdf/1901.07291.pdf (section 3.1)
    :dataset_size: A dictionary contains the size (value) of each of the language (key).
    """
    tot_size = 0
    probs = OrderedDict()
    for lang, size in dataset_size.items():
        tot_size += size
    for lang, size in dataset_size.items():
        probs[lang] = size / tot_size

    pen_prob = OrderedDict()
    tot_pen_prob = 0.0
    for lang, prob in probs.items():
        tot_pen_prob += (prob ** alpha)
    sum_ = 0.0
    for lang, prob in probs.items():
        pen_prob[lang] = (prob ** alpha) / tot_pen_prob
        sum_ += pen_prob[lang]
    assert math.fabs(1 - sum_) < 1e-6
    return pen_prob