"""
Trainer module tests
"""

import os
import unittest
import tempfile

import torch

from transformers import AutoTokenizer, AutoModelForSequenceClassification

from txtai.data import Data
from txtai.pipeline import HFTrainer, Labels, Questions, Sequences


class TestTrainer(unittest.TestCase):
    """
    Trainer tests.
    """

    @classmethod
    def setUpClass(cls):
        """
        Create default datasets.
        """

        cls.data = [{"text": "Dogs", "label": 0}, {"text": "dog", "label": 0}, {"text": "Cats", "label": 1}, {"text": "cat", "label": 1}] * 100

    def testBasic(self):
        """
        Test training a model with basic parameters
        """

        trainer = HFTrainer()
        model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", self.data)

        labels = Labels((model, tokenizer), dynamic=False)
        self.assertEqual(labels("cat")[0][0], 1)

    def testCLM(self):
        """
        Test training a model with causal language modeling
        """

        trainer = HFTrainer()

        # Test default parameters
        model, _ = trainer("hf-internal-testing/tiny-random-gpt2", self.data, maxlength=16, task="language-generation")
        self.assertIsNotNone(model)

        # Test pack merging
        model, _ = trainer("hf-internal-testing/tiny-random-gpt2", self.data, maxlength=16, task="language-generation", merge="pack")
        self.assertIsNotNone(model)

        # Test no merging
        model, _ = trainer("hf-internal-testing/tiny-random-gpt2", self.data, maxlength=16, task="language-generation", merge=None)
        self.assertIsNotNone(model)

    def testCustom(self):
        """
        Test training a model with custom parameters
        """

        # pylint: disable=E1120
        model = AutoModelForSequenceClassification.from_pretrained("google/bert_uncased_L-2_H-128_A-2")
        tokenizer = AutoTokenizer.from_pretrained("google/bert_uncased_L-2_H-128_A-2")

        trainer = HFTrainer()
        model, tokenizer = trainer(
            (model, tokenizer),
            self.data,
            self.data,
            columns=("text", "label"),
            do_eval=True,
            output_dir=os.path.join(tempfile.gettempdir(), "trainer"),
        )

        labels = Labels((model, tokenizer), dynamic=False)
        self.assertEqual(labels("cat")[0][0], 1)

    def testDataFrame(self):
        """
        Test training a model with a mock pandas DataFrame
        """

        class TestDataFrame:
            """
            Test DataFrame
            """

            def __init__(self, data):
                # Get list of columns
                self.columns = list(data[0].keys())

                # Build columnar data view
                self.data = {}
                for column in self.columns:
                    self.data[column] = Values([row[column] for row in data])

            def __getitem__(self, column):
                return self.data[column]

        class Values:
            """
            Test values list
            """

            def __init__(self, values):
                self.values = list(values)

            def __getitem__(self, index):
                return self.values[index]

            def unique(self):
                """
                Returns a list of unique values.

                Returns:
                    unique list of values
                """

                return set(self.values)

        # Mock DataFrame
        df = TestDataFrame(self.data)

        trainer = HFTrainer()
        model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", df)

        labels = Labels((model, tokenizer), dynamic=False)
        self.assertEqual(labels("cat")[0][0], 1)

    def testDataset(self):
        """
        Test training a model with a mock Hugging Face Dataset
        """

        class TestDataset(torch.utils.data.Dataset):
            """
            Test Dataset
            """

            def __init__(self, data):
                self.data = data
                self.unique = lambda _: [0, 1]

            def __len__(self):
                return len(self.data)

            def __getitem__(self, index):
                return self.data[index]

            def column_names(self):
                """
                Returns column names for this dataset

                Returns:
                    list of columns
                """

                return ["text", "label"]

            # pylint: disable=W0613
            def map(self, fn, batched, num_proc, remove_columns):
                """
                Map each dataset row using fn.

                Args:
                    fn: function
                    batched: batch records

                Returns:
                    updated Dataset
                """

                self.data = [fn(x) for x in self.data]
                return self

        ds = TestDataset(self.data)

        trainer = HFTrainer()
        model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", ds)

        labels = Labels((model, tokenizer), dynamic=False)
        self.assertEqual(labels("cat")[0][0], 1)

    def testEmpty(self):
        """
        Test an empty training data object
        """

        self.assertIsNone(Data(None, None, None).process(None))

    def testKD(self):
        """
        Test knowledge distillation
        """

        # Base model
        trainer = HFTrainer()
        model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", self.data)

        # Train with knowledge distillation
        model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", self.data, teacher=(model, tokenizer))

        labels = Labels((model, tokenizer), dynamic=False)
        self.assertEqual(labels("cat")[0][0], 1)

    def testMLM(self):
        """
        Test training a model with masked language modeling.
        """

        trainer = HFTrainer()
        model, _ = trainer("hf-internal-testing/tiny-random-bert", self.data, task="language-modeling")

        # Test model completed successfully
        self.assertIsNotNone(model)

    def testMultiLabel(self):
        """
        Test training model with labels provided as a list
        """

        data = []
        for x in self.data:
            data.append({"text": x["text"], "label": [0.0, 1.0] if x["label"] else [1.0, 0.0]})

        trainer = HFTrainer()
        model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", data)

        labels = Labels((model, tokenizer), dynamic=False)
        self.assertEqual(labels("cat")[0][0], 1)

    def testPEFT(self):
        """
        Test training a model with causal language modeling and PEFT
        """

        trainer = HFTrainer()
        model, _ = trainer(
            "hf-internal-testing/tiny-random-gpt2",
            self.data,
            maxlength=16,
            task="language-generation",
            quantize=True,
            lora=True,
        )

        # Test model completed successfully
        self.assertIsNotNone(model)

    def testQA(self):
        """
        Test training a QA model
        """

        # Training data
        data = [
            {"question": "What ingredient?", "context": "1 can whole tomatoes", "answers": "tomatoes"},
            {"question": "What ingredient?", "context": "Crush 1 tomato", "answers": "tomato"},
            {"question": "What ingredient?", "context": "1 yellow onion", "answers": "onion"},
            {"question": "What ingredient?", "context": "Unwrap 2 red onions", "answers": "onions"},
            {"question": "What ingredient?", "context": "1 red pepper", "answers": "pepper"},
            {"question": "What ingredient?", "context": "Clean 3 red peppers", "answers": "peppers"},
            {"question": "What ingredient?", "context": "1 clove garlic", "answers": "garlic"},
            {"question": "What ingredient?", "context": "Unwrap 3 cloves of garlic", "answers": "garlic"},
            {"question": "What ingredient?", "context": "3 pieces of ginger", "answers": "ginger"},
            {"question": "What ingredient?", "context": "Peel 1 orange", "answers": "orange"},
            {"question": "What ingredient?", "context": "1/2 lb beef", "answers": "beef"},
            {"question": "What ingredient?", "context": "Roast 3 lbs of beef", "answers": "beef"},
            {"question": "What ingredient?", "context": "1 pack of chicken", "answers": "chicken"},
            {"question": "What ingredient?", "context": "Forest through the trees", "answers": None},
        ]

        trainer = HFTrainer()
        model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", data, data, task="question-answering", num_train_epochs=40)

        questions = Questions((model, tokenizer), gpu=True)
        self.assertTrue("onion" in questions(["What ingredient?"], ["Peel 1 onion"])[0])

    def testRegression(self):
        """
        Test training a model with a regression (continuous) output
        """

        data = []
        for x in self.data:
            data.append({"text": x["text"], "label": x["label"] + 0.1})

        trainer = HFTrainer()
        model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", data)

        labels = Labels((model, tokenizer), dynamic=False)

        # Regression tasks return a single entry with the regression output
        self.assertGreater(labels("cat")[0][1], 0.5)

    def testRTD(self):
        """
        Test training a language model with replaced token detection
        """

        # Save directory
        output = os.path.join(tempfile.gettempdir(), "trainer.rtd")

        trainer = HFTrainer()
        model, _ = trainer("hf-internal-testing/tiny-random-electra", self.data, task="token-detection", output_dir=output)

        # Test model completed successfully
        self.assertIsNotNone(model)

        # Test output directories exist
        self.assertTrue(os.path.exists(os.path.join(output, "generator")))
        self.assertTrue(os.path.exists(os.path.join(output, "discriminator")))

    def testSeqSeq(self):
        """
        Test training a sequence-sequence model
        """

        data = [
            {"source": "Running again", "target": "Sleeping again"},
            {"source": "Run", "target": "Sleep"},
            {"source": "running", "target": "sleeping"},
        ]

        trainer = HFTrainer()
        model, tokenizer = trainer("t5-small", data, task="sequence-sequence", prefix="translate Run to Sleep: ", learning_rate=1e-3)

        # Run run-sleep translation
        sequences = Sequences((model, tokenizer))
        result = sequences("translate Run to Sleep: run")
        self.assertEqual(result.lower(), "sleep")