Lateinicus/scripts/generate_test_levels.py

import pymongo
import markovify
import sys
import random
import os

def log(msg, err=False, tabs=0):
    if (not err):
        print("[*] " + "\t" * tabs + msg)
    else:
        print("[X] " + "\t" * tabs + msg)

def dbg(msg, tabs=0):
    print("[D] " + "\t" * tabs + msg)

if (len(sys.argv) < 3):
    log("Not enough arguments!", err=True)
    sys.exit(1)

# Load the model for the markov chains
file = open("bee-movie.txt", "r")
text = file.read()
file.close()

# Create the model using Markov Chains
model = markovify.Text(text)

# Generate 25 levels
log("Generating levels...")
levels = []
vocab = list(range(1, 125))
level = 1
for i in range(25):
    name = model.make_short_sentence(50, tries=100)
    description = " ".join([model.make_sentence(tries=100) for x in range(4)])
    level_vocab = []

    # Get random vocabulary
    for v in range(4):
        index = random.randint(0, len(vocab) - 1)
        level_vocab.append(index)
        del vocab[index]
        
    log("Level {}: {}".format(level, name))
    levels.append({
        "level": level,
        "name": name,
        "description": description,
        "vocab": level_vocab
    })
    level += 1

if os.getenv("DEBUG") != None:
    sys.exit(0)
    
log("Connecting to database...")    
client = pymongo.MongoClient(sys.argv[1])
log("Getting DB...")
db = client[sys.argv[2]]
log("Inserting...")
res = db["levels"].insert_many(levels)

log("Success", tabs=1)