#!/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.
"""
Test Teachers.

A module for testing various teacher types in ParlAI
"""

import os
import unittest
from parlai.utils import testing as testing_utils
from parlai.core.teachers import DialogTeacher, create_task_agent_from_taskname
from parlai.core.metrics import SumMetric
import regex as re
from parlai.core.message import Message
from parlai.core.opt import Opt
import parlai.utils.logging as logging
from parlai.utils.io import PathManager


class TestAbstractImageTeacher(unittest.TestCase):
    """
    Test AbstractImageTeacher.
    """

    def _test_display_output(self, image_mode):
        """
        Test display data output with given image_mode.
        """
        with testing_utils.tempdir() as tmpdir:
            data_path = tmpdir
            PathManager.mkdirs(os.path.join(data_path, 'ImageTeacher'))

            opt = {
                'task': 'integration_tests:ImageTeacher',
                'datapath': data_path,
                'image_mode': image_mode,
                'verbose': True,
            }
            output = testing_utils.display_data(opt)
            train_labels = re.findall(r"\[labels\].*\n", output[0])
            valid_labels = re.findall(r"\[eval_labels\].*\n", output[1])
            test_labels = re.findall(r"\[eval_labels\].*\n", output[2])

            for i, lbls in enumerate([train_labels, valid_labels, test_labels]):
                self.assertGreater(len(lbls), 0, 'DisplayData failed')
                self.assertEqual(len(lbls), len(set(lbls)), output[i])

    def test_display_data_no_image(self):
        """
        Test that, with no images loaded, all examples are different.
        """
        self._test_display_output('no_image_model')

    @testing_utils.skipUnlessVision
    @testing_utils.skipUnlessGPU
    def test_display_data_resnet(self):
        """
        Test that, with pre-loaded image features, all examples are different.
        """
        self._test_display_output('resnet152')


class TestParlAIDialogTeacher(unittest.TestCase):
    def test_good_fileformat(self):
        """
        Checks that we fail to load a dataset where the use specified eval_labels.
        """
        with testing_utils.tempdir() as tmpdir:
            fp = os.path.join(tmpdir, "goodfile.txt")
            with PathManager.open(fp, "w") as f:
                f.write('id:test_file\ttext:input\tlabels:good label\n\n')
            opt = {'task': 'fromfile', 'fromfile_datapath': fp, 'verbose': True}
            testing_utils.display_data(opt)

    def test_bad_fileformat(self):
        """
        Checks that we fail to load a dataset where the use specified eval_labels.
        """
        with testing_utils.tempdir() as tmpdir:
            fp = os.path.join(tmpdir, "badfile.txt")
            with PathManager.open(fp, "w") as f:
                f.write('id:test_file\ttext:input\teval_labels:bad label\n\n')
            opt = {'task': 'fromfile', 'fromfile_datapath': fp, 'verbose': True}
            with self.assertRaises(ValueError):
                testing_utils.display_data(opt)

    def test_no_text(self):
        with testing_utils.tempdir() as tmpdir:
            fp = os.path.join(tmpdir, "badfile.txt")
            with PathManager.open(fp, "w") as f:
                f.write('id:test_file\tlabels:bad label\n\n')
            opt = {'task': 'fromfile', 'fromfile_datapath': fp, 'verbose': True}
            with self.assertRaises(ValueError):
                testing_utils.display_data(opt)

    def test_no_labels(self):
        with testing_utils.tempdir() as tmpdir:
            fp = os.path.join(tmpdir, "badfile.txt")
            with PathManager.open(fp, "w") as f:
                f.write('id:test_file\ttext:bad text\n\n')
            opt = {'task': 'fromfile', 'fromfile_datapath': fp, 'verbose': True}
            with self.assertRaises(ValueError):
                testing_utils.display_data(opt)

    def test_one_episode(self):
        with testing_utils.tempdir() as tmpdir:
            fp = os.path.join(tmpdir, "badfile.txt")
            with PathManager.open(fp, "w") as f:
                for _ in range(1000):
                    f.write('id:test_file\ttext:placeholder\tlabels:placeholder\n\n')
            opt = {'task': 'fromfile', 'fromfile_datapath': fp, 'verbose': True}
            with self.assertLogs(logger=logging.logger, level='DEBUG') as cm:
                testing_utils.display_data(opt)
                print("\n".join(cm.output))
                assert any('long episode' in l for l in cm.output)

            # invert the logic of the assertion
            with self.assertRaises(self.failureException):
                fp = os.path.join(tmpdir, "goodfile.txt")
                with PathManager.open(fp, "w") as f:
                    for _ in range(1000):
                        f.write(
                            'id:test_file\ttext:placeholder\tlabels:placeholder\tepisode_done:True\n\n'
                        )
                opt = {'task': 'fromfile', 'fromfile_datapath': fp, 'verbose': True}
                with self.assertLogs(logger=logging.logger, level='DEBUG') as cm:
                    testing_utils.display_data(opt)
                    assert any('long episode' in l for l in cm.output)


class TestConversationTeacher(unittest.TestCase):
    def test_good_fileformat(self):
        """
        Checks that we succeed in loading a well formatted jsonl file.
        """
        with testing_utils.tempdir() as tmpdir:
            fp = os.path.join(tmpdir, "goodfile.jsonl")
            with PathManager.open(fp, "w") as f:
                f.write(
                    '{"dialog": [[{"text": "Hi.", "id": "speaker1"}, {"text": "Hello.", "id": "speaker2"}]]}\n'
                )
            opt = {'task': 'jsonfile', 'jsonfile_datapath': fp, 'verbose': True}
            testing_utils.display_data(opt)

    def test_no_text(self):
        with testing_utils.tempdir() as tmpdir:
            fp = os.path.join(tmpdir, "badfile.jsonl")
            with PathManager.open(fp, "w") as f:
                f.write(
                    '{"dialog": [[{"id": "speaker1"}, {"text": "Hello.", "id": "speaker2"}]]}\n'
                )
            opt = {'task': 'jsonfile', 'jsonfile_datapath': fp, 'verbose': True}

            train_out, _, _ = testing_utils.display_data(opt)
            self.assertTrue("'text': None" in train_out.split('\n')[1])

    def test_firstspeaker_label(self):
        with testing_utils.tempdir() as tmpdir:
            fp = os.path.join(tmpdir, "goodfile.jsonl")
            with PathManager.open(fp, "w") as f:
                f.write(
                    '{"dialog": [[{"text": "Hi.", "id": "speaker1"}, {"text": "Hello.", "id": "speaker2"}]]}\n'
                )
            opt = {
                'task': 'jsonfile',
                'jsonfile_datapath': fp,
                'verbose': True,
                'label_turns': 'firstspeaker',
            }
            train_out, valid_out, test_out = testing_utils.display_data(opt)
            texts = [
                l.split(':', 1)[-1].strip()
                for l in train_out.split('\n')
                if l in train_out
                if '[text]' in l
            ]
            labels = [
                l.split(':', 1)[-1].strip()
                for l in train_out.split('\n')
                if l in train_out
                if 'labels' in l
            ]
            self.assertEqual(texts[0], '__SILENCE__')
            self.assertEqual(labels[0], 'Hi.')

    def test_secondspeaker_label(self):
        with testing_utils.tempdir() as tmpdir:
            fp = os.path.join(tmpdir, "goodfile.jsonl")
            with PathManager.open(fp, "w") as f:
                f.write(
                    '{"dialog": [[{"text": "Hi.", "id": "speaker1"}, {"text": "Hello.", "id": "speaker2"}]]}\n'
                )
            opt = {
                'task': 'jsonfile',
                'jsonfile_datapath': fp,
                'verbose': True,
                'label_turns': 'secondspeaker',
            }
            train_out, valid_out, test_out = testing_utils.display_data(opt)
            texts = [
                l.split(':', 1)[-1].strip()
                for l in train_out.split('\n')
                if l in train_out
                if '[text]' in l
            ]
            labels = [
                l.split(':', 1)[-1].strip()
                for l in train_out.split('\n')
                if l in train_out
                if 'labels' in l
            ]
            self.assertEqual(texts[0], 'Hi.')
            self.assertEqual(labels[0], 'Hello.')

    def test_both_label(self):
        with testing_utils.tempdir() as tmpdir:
            fp = os.path.join(tmpdir, "goodfile.jsonl")
            with PathManager.open(fp, "w") as f:
                f.write(
                    '{"dialog": [[{"text": "Hi.", "id": "speaker1"}, {"text": "Hello.", "id": "speaker2"}]]}\n'
                )
            opt = {
                'task': 'jsonfile',
                'jsonfile_datapath': fp,
                'verbose': True,
                'label_turns': 'both',
            }
            train_out, valid_out, test_out = testing_utils.display_data(opt)
            texts = [
                l.split(':', 1)[-1].strip()
                for l in train_out.split('\n')
                if l in train_out
                if '[text]' in l
            ]
            labels = [
                l.split(':', 1)[-1].strip()
                for l in train_out.split('\n')
                if l in train_out
                if 'labels' in l
            ]
            num_episodes = train_out.count("END OF EPISODE")
            self.assertEqual(texts[0], '__SILENCE__')
            self.assertEqual(labels[0], 'Hi.')
            self.assertEqual(texts[1], 'Hi.')
            self.assertEqual(labels[1], 'Hello.')
            self.assertEqual(num_episodes, 2)


class TestChunkTeacher(unittest.TestCase):
    """
    Test chunked teacher.
    """

    def test_iter(self):
        with testing_utils.tempdir() as tmpdir:
            teacher = create_task_agent_from_taskname(
                {'task': 'integration_tests', 'datatype': 'valid', 'datapath': tmpdir}
            )[0]
            # twice to assert we reset iterators correctly
            assert len(list(teacher)) == 100
            assert len(list(teacher)) == 100

    def test_no_batched(self):
        valid, test = testing_utils.eval_model(
            dict(task='integration_tests:chunky', model='repeat_label'),
            valid_datatype='valid:stream',
            test_datatype='test:stream',
        )
        assert valid['exs'] == 100
        assert test['exs'] == 100

    def test_batched(self):
        valid, test = testing_utils.eval_model(
            dict(
                task='integration_tests:chunky',
                model='parlai.agents.test_agents.test_agents:MockTorchAgent',
                batchsize=32,
            ),
            valid_datatype='valid:stream',
            test_datatype='test:stream',
        )
        assert valid['exs'] == 100
        assert test['exs'] == 100

    def test_dynamic_batched(self):
        valid, test = testing_utils.eval_model(
            dict(
                task='integration_tests:chunky',
                model='parlai.agents.test_agents.test_agents:MockTorchAgent',
                datatype='valid:stream',
                batchsize=32,
                truncate=16,
                dynamic_batching='full',
            ),
            valid_datatype='valid:stream',
            test_datatype='test:stream',
        )
        assert valid['exs'] == 100
        assert test['exs'] == 100

    def test_non_stream_works(self):
        testing_utils.eval_model(
            dict(
                task='integration_tests:chunky',
                model='parlai.agents.test_agents.test_agents:MockTorchAgent',
                batchsize=32,
            ),
            valid_datatype='valid',
        )

        testing_utils.eval_model(
            dict(
                task='integration_tests:chunky',
                model='parlai.agents.test_agents.test_agents:MockTorchAgent',
                batchsize=32,
            ),
            valid_datatype='valid:stream',
            test_datatype='test',
        )


class CustomEvaluationTeacher(DialogTeacher):
    def __init__(self, opt, shared=None):
        opt['datafile'] = 'mock'
        super().__init__(opt, shared)

    def custom_evaluation(self, teacher_action, label, model_response):
        self.metrics.add('contains1', SumMetric(int('1' in model_response['text'])))

    def setup_data(self, fold):
        yield ('1 2', '1 2'), True
        yield ('3 4', '3 4'), True


class TestCustomEvaluation(unittest.TestCase):
    def test_custom_eval(self):
        opt = {'task': 'custom', 'datatype': 'valid'}
        teacher = CustomEvaluationTeacher(opt)
        teacher.act()
        teacher.observe({'text': 'a b'})
        teacher.act()
        teacher.observe({'text': '1 2'})
        report = teacher.report()
        assert 'contains1' in report
        assert report['contains1'] == 1
        assert report['exs'] == 2


class _MockTeacher(DialogTeacher):
    def __init__(self, opt, shared=None):
        opt['datafile'] = 'mock'
        super().__init__(opt)


class TupleTeacher(_MockTeacher):
    def setup_data(self, datafile):
        for _ in range(3):
            for j in range(1, 4):
                yield (str(j), str(j * 2)), j == 1


class DictTeacher(_MockTeacher):
    def setup_data(self, datafile):
        for _ in range(3):
            for j in range(1, 4):
                yield {'text': str(j), 'label': str(j * 2)}, j == 1


class MessageTeacher(_MockTeacher):
    def setup_data(self, datafile):
        for _ in range(3):
            for j in range(1, 4):
                yield Message({'text': str(j), 'label': str(j * 2)}), j == 1


class NoDatafileTeacher(DialogTeacher):
    def setup_data(self, datafile):
        yield Message({'text': datafile, 'label': datafile}), True


class ViolationTeacher(_MockTeacher):
    def setup_data(self, datafile):
        yield {'text': 'foo', 'episode_done': True}, True


class TestDialogTeacher(unittest.TestCase):
    def test_iter(self):
        opt = Opt({'datatype': 'valid', 'datapath': '/tmp', 'task': 'test'})
        teacher = TupleTeacher(opt)
        # twice to ensure we reset iterators correctly
        examples = list(teacher)
        assert len(examples) == 9
        examples = list(teacher)
        assert len(examples) == 9

    def test_nodatafile(self):
        for dt in [
            'train:ordered',
            'train:stream:ordered',
            'valid',
            'test',
            'valid:stream',
            'test:stream',
        ]:
            opt = Opt({'datatype': dt, 'datapath': '/tmp', 'task': 'test'})
            with self.assertRaises(KeyError):
                NoDatafileTeacher(opt)

    def _verify_act(self, act, goal_text, goal_label, episode_done):
        assert 'eval_labels' in act or 'labels' in act
        labels = act.get('labels', act.get('eval_labels'))
        assert isinstance(labels, tuple)
        assert len(labels) == 1
        assert act['text'] == str(goal_text)
        assert labels[0] == str(goal_label)

    def _test_iterate(self, teacher_class):
        for dt in [
            'train:ordered',
            'train:stream:ordered',
            'valid',
            'test',
            'valid:stream',
            'test:stream',
        ]:
            opt = Opt({'datatype': dt, 'datapath': '/tmp', 'task': 'test'})
            teacher = teacher_class(opt)

            self._verify_act(teacher.act(), 1, 2, False)
            self._verify_act(teacher.act(), 2, 4, False)
            self._verify_act(teacher.act(), 3, 6, True)

            self._verify_act(teacher.act(), 1, 2, False)
            self._verify_act(teacher.act(), 2, 4, False)
            self._verify_act(teacher.act(), 3, 6, True)

            self._verify_act(teacher.act(), 1, 2, False)
            self._verify_act(teacher.act(), 2, 4, False)
            self._verify_act(teacher.act(), 3, 6, True)

            assert teacher.epoch_done()

    def test_tuple_teacher(self):
        self._test_iterate(TupleTeacher)

    def test_dict_teacher(self):
        self._test_iterate(DictTeacher)

    def test_message_teacher(self):
        self._test_iterate(MessageTeacher)

    def test_violation_teacher(self):
        with self.assertRaises(KeyError):
            self._test_iterate(ViolationTeacher)


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