# The basic version of "new poem" generator.

from random import randint, choice

def wordCountSum(wordCount):
    sum = 0
    for value in wordCount.values():
        sum += value
    return sum

def retrieveRandomWord(wordCount):

    randIndex = randint(1, wordCountSum(wordCount))
    for word, value in wordCount.items():
        randIndex -= value
        if randIndex <= 0:
            return word

def buildWordDict(fn):
    d = {}
    infile =  open(fn, "r")
    body = infile.read()
    infile.close()
    addWord(body, d)
    return d

def addWord(words, wordDict):
    words = words.replace("？", "")

    for i in range(1, len(words)):
        if words[i-1] not in wordDict:
            #Create a new dictionary for this word
            wordDict[words[i-1]] = {}
        if words[i] not in wordDict[words[i-1]]:
            wordDict[words[i-1]][words[i]] = 0
        wordDict[words[i-1]][words[i]] += 1

    return wordDict

def main():
    fn = 'gdrive/MyDrive/train.txt'
    wordDict = buildWordDict(fn)

    #Generate text of length 100
    MAX_LEN = 100
    initialWord = choice(list(wordDict.keys()))
    text = ""
    currentWord = initialWord
    i = 0
    while True:
        i += 1
        text += currentWord
        if currentWord not in wordDict:
            currentWord = initialWord
        currentWord = retrieveRandomWord(wordDict[currentWord])
        if currentWord == '\n' and i >= MAX_LEN: break

    print(text)

main()