from datetime import datetime

from glossolalia.loader import load_seeds, load_text
from glossolalia.lstm import LisSansTaMaman

"""
he a mighty hand
"""


def train():
    # should_train = True
    nb_words = 60
    nb_epoch = 60
    nb_layers = 128
    dropout = .3
    validation_split = .3
    lstm = LisSansTaMaman(nb_layers, dropout, validation_split, debug=True)
    filename_output = "./output/verbe_%i-d%.1f_%s.txt" % (
        nb_layers, dropout, datetime.now().strftime("%y%m%d_%H%M"))
    seeds = ["Thus", "And", "He", "Our", "I", "Thou", "But"]

    corpus = load_text("./data.txt")
    print("Corpus:", corpus[:10])
    lstm.create_model(corpus[:100])  # FIXME: Scale up to whole corpus

    with open(filename_output, "a+") as f:
        for i in range(0, nb_epoch, 10):
            lstm.fit(epochs=min(i + 10, nb_epoch), initial_epoch=i,
                     validation_split=validation_split)

            for output in lstm.predict_seeds(nb_words, seeds):
                print(output)
                f.writelines(output)

        for i, seed in enumerate(load_seeds(corpus, 5)):
            output = lstm.predict(seed, nb_words)
            print("%i %s -> %s" % (i, seed, output))
            f.writelines(output)

        while True:
            input_text = input("> ")
            text = lstm.predict(input_text, nb_words)
            print(text)
            f.writelines("%s\n" % text)


if __name__ == '__main__':
    train()
