#!/usr/bin/env python3

# Copyright (c) Facebook, Inc. and its affiliates.
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from parlai.core.agents import create_agent
import unittest
import parlai.scripts.display_data as display_data
from parlai.tasks.wizard_of_wikipedia.agents import TOKEN_KNOWLEDGE
import parlai.utils.testing as testing_utils

from projects.wizard_of_wikipedia.knowledge_retriever.knowledge_retriever import (
    KnowledgeRetrieverAgent,
)

END2END_OPTIONS = {
    'task': 'wizard_of_wikipedia:generator:random_split',
    'model_file': 'zoo:wizard_of_wikipedia/end2end_generator/model',
    'batchsize': 4,
    'log_every_n_secs': 30,
    'embedding_type': 'random',
}


RETRIEVAL_OPTIONS = {
    'task': 'wizard_of_wikipedia',
    'model': 'projects:wizard_of_wikipedia:wizard_transformer_ranker',
    'model_file': 'zoo:wizard_of_wikipedia/full_dialogue_retrieval_model/model',
    'datatype': 'test',
    'n_heads': 6,
    'batchsize': 4,
    'ffn_size': 1200,
    'embeddings_scale': False,
    'delimiter': ' __SOC__ ',
    'n_positions': 1000,
    'legacy': True,
}


@testing_utils.skipUnlessGPU
class TestWizardModel(unittest.TestCase):
    """
    Checks that pre-trained Wizard models give the correct results.
    """

    @classmethod
    def setUpClass(cls):
        # go ahead and download things here
        parser = display_data.setup_args()
        parser.set_defaults(**END2END_OPTIONS)
        opt = parser.parse_args([])
        opt['num_examples'] = 1
        opt['verbose'] = True
        display_data.display_data(opt)

    def test_end2end(self):
        valid, _ = testing_utils.eval_model(END2END_OPTIONS, skip_test=True)
        self.assertAlmostEqual(valid['ppl'], 61.21, places=2)
        self.assertAlmostEqual(valid['f1'], 0.1717, places=4)
        self.assertAlmostEqual(valid['know_acc'], 0.2201, places=4)

    def test_retrieval(self):
        _, test = testing_utils.eval_model(RETRIEVAL_OPTIONS, skip_valid=True)
        self.assertAlmostEqual(test['accuracy'], 0.8631, places=4)
        self.assertAlmostEqual(test['hits@5'], 0.9814, places=4)
        self.assertAlmostEqual(test['hits@10'], 0.9917, places=4)


class TestKnowledgeRetriever(unittest.TestCase):
    """
    Checks that the Knowledge Retriever module returns the correct results.
    """

    def test_knowledge_retriever(self):
        from parlai.core.params import ParlaiParser

        parser = ParlaiParser(False, False)
        KnowledgeRetrieverAgent.add_cmdline_args(parser, partial_opt=None)
        parser.set_params(
            model='projects:wizard_of_wikipedia:knowledge_retriever',
            add_token_knowledge=True,
        )
        knowledge_opt = parser.parse_args([])
        knowledge_agent = create_agent(knowledge_opt)

        knowledge_agent.observe(
            {
                'text': 'what do you think of mountain dew?',
                'chosen_topic': 'Mountain Dew',
                'episode_done': False,
            }
        )

        knowledge_act = knowledge_agent.act()

        title = knowledge_act['title']
        self.assertEqual(title, 'Mountain Dew', 'Did not save chosen topic correctly')

        knowledge = knowledge_act['text']
        self.assertIn(
            TOKEN_KNOWLEDGE, knowledge, 'Knowledge token was not inserted correctly'
        )

        checked_sentence = knowledge_act['checked_sentence']
        self.assertEqual(
            checked_sentence,
            'Mountain Dew (stylized as Mtn Dew) is a carbonated soft drink brand produced and owned by PepsiCo.',
            'Did not correctly choose the checked sentence',
        )


if __name__ == '__main__':
    unittest.main()
