import random
import json
import argparse

def select(model: dict):
    m_filtered = list(filter(lambda i: i[0] != "_count", model.items()))
    return random.choices(
        list(map(lambda i: i[0], m_filtered)),
        list(map(lambda i: i[1]["_count"], m_filtered)),
        k=1
    )[0]

def generate_token(model, former, current):
    p = random.random()
    if p <= 0.90:
        try:
            return select(model[former][current])
        except: pass
    if p <= 0.98:
        try:
            return select(model[current])
        except: pass
    return select(model)

def main(args):
    with open(args.model) as m_file: model = json.load(m_file)
    former = select(model)
    current = select(model[former])
    print(f"{former} {current} ", end="")
    for _ in range(args.size - 2):
        prior = generate_token(model, former, current)
        print(prior + " ", end="")
        former = current
        current = prior


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('model')
    parser.add_argument('-s', '--size', default=100, type=int)
    main(parser.parse_args())