import collections

corpus = {
    "l o w </w>": 5,
    "l o w e r </w>": 2,
    "n e w e s t </w>": 6,
    "w i d e s t </w>": 3,
    "の ら 猫 </w>": 4
}

def get_stats(corpus):
    pairs = collections.defaultdict(int)
    for word, freq in corpus.items():
        symbols = word.split()
        for i in range(len(symbols) - 1):
            pairs[(symbols[i], symbols[i + 1])] += freq
    return pairs

def merge_vocab(pair, corpus):
    new_corpus = {}
    bigram = " ".join(pair)
    replacement = "".join(pair)
    for word, freq in corpus.items():
        new_word = word.replace(bigram, replacement)
        new_corpus[new_word] = freq
    return new_corpus

num_merges = 10
for i in range(num_merges):
    pairs = get_stats(corpus)
    best = max(pairs, key=pairs.get)
    corpus = merge_vocab(best, corpus)
    print(f"Merge #{i+1}: {best[0]} + {best[1]}")

for word in corpus:
    print(word)

# 佐藤 聖
# Attention Is All You Need (Vaswani et al., 2017)
# Related to data preprocessing used before Section 3.4
# (Embeddings and Softmax) in Vaswani et al. (2017)
# SubwordBPE.py
# 第2部 第0章