import random


def random_word2(tokens, vocab_range, mask, mlm_positions, position_to_prob):
    """
    Masking some random tokens for Language Model task with probabilities as in
        the original BERT paper.
    :param tokens: list of int, tokenized sentence.
    :param vocab_range: for choosing a random word
    :param mlm_positions: the positions sequence. if select, all tokens in the positions would be masked.
    :param position_to_prob: the sampling probability of each token in the specific position
    :return: (list of int, list of int), masked tokens and related labels for
        LM prediction
    """
    random.shuffle(mlm_positions)
    output_label = [-1] * len(tokens)
    for positions in mlm_positions:
        sample_prob = sum([position_to_prob.get(position, 0.15) for position in positions]) / max(len(positions), 1)
        prob = random.random()
        if prob < sample_prob:
            prob /= sample_prob

            for position in positions:
                token = tokens[position]
                if output_label[position] == -1:
                    prob2 = random.random()
                    if prob2 < 0.8:
                        tokens[position] = mask
                    elif prob2 < 0.9:
                        tokens[position] = random.choice(list(range(*vocab_range)))
                    output_label[position] = token

    if all(o == -1 for o in output_label):
        # at least mask 1
        select_positions = mlm_positions[0]
        for position in select_positions:
            token = tokens[position]
            prob2 = random.random()
            if prob2 < 0.8:
                tokens[position] = mask
            elif prob2 < 0.9:
                tokens[position] = random.choice(list(range(*vocab_range)))
            output_label[position] = token

    return tokens, output_label