import shutil

from flair.data import Dictionary
from flair.trainers.language_model_trainer import TextCorpus


def test_train_resume_language_model_training(resources_path, results_base_path, tasks_base_path):
    # get default dictionary
    dictionary: Dictionary = Dictionary.load('chars')

    # get the example corpus and process at character level in forward direction
    corpus: TextCorpus = TextCorpus(resources_path / 'corpora/lorem_ipsum',
                                dictionary,
                                forward=True,
                                character_level=True)

    assert (corpus.test is not None)
    assert (corpus.train is not None)
    assert (corpus.valid is not None)
    assert (len(corpus.train) == 2)


def test_generate_text_with_small_temperatures():

    from flair.embeddings import FlairEmbeddings

    language_model = FlairEmbeddings('news-forward-fast').lm

    text, likelihood = language_model.generate_text(temperature=0.01, number_of_characters=100)
    assert (text is not None)
    assert (len(text) >= 100)


def test_compute_perplexity():

    from flair.embeddings import FlairEmbeddings

    language_model = FlairEmbeddings('news-forward-fast').lm

    grammatical = 'The company made a profit'
    perplexity_gramamtical_sentence = language_model.calculate_perplexity(grammatical)

    ungrammatical = 'Nook negh qapla!'
    perplexity_ungramamtical_sentence = language_model.calculate_perplexity(ungrammatical)

    print(f'"{grammatical}" - perplexity is {perplexity_gramamtical_sentence}')
    print(f'"{ungrammatical}" - perplexity is {perplexity_ungramamtical_sentence}')

    assert (perplexity_gramamtical_sentence < perplexity_ungramamtical_sentence)

    language_model = FlairEmbeddings('news-backward-fast').lm

    grammatical = 'The company made a profit'
    perplexity_gramamtical_sentence = language_model.calculate_perplexity(grammatical)

    ungrammatical = 'Nook negh qapla!'
    perplexity_ungramamtical_sentence = language_model.calculate_perplexity(ungrammatical)

    print(f'"{grammatical}" - perplexity is {perplexity_gramamtical_sentence}')
    print(f'"{ungrammatical}" - perplexity is {perplexity_ungramamtical_sentence}')

    assert (perplexity_gramamtical_sentence < perplexity_ungramamtical_sentence)
