#!/usr/bin/env python3

import argparse
import os
import unittest
from inspect import currentframe, getframeinfo

import numpy as np
import torch
from examples.speech_recognition.data.data_utils import lengths_to_encoder_padding_mask
from fairseq.data import data_utils as fairseq_data_utils
from fairseq.data.dictionary import Dictionary
from fairseq.models import (
    BaseFairseqModel,
    FairseqDecoder,
    FairseqEncoder,
    FairseqEncoderDecoderModel,
    FairseqEncoderModel,
    FairseqModel,
)
from fairseq.tasks.fairseq_task import LegacyFairseqTask


DEFAULT_TEST_VOCAB_SIZE = 100


# ///////////////////////////////////////////////////////////////////////////
# utility function to setup dummy dict/task/input
# ///////////////////////////////////////////////////////////////////////////


def get_dummy_dictionary(vocab_size=DEFAULT_TEST_VOCAB_SIZE):
    dummy_dict = Dictionary()
    # add dummy symbol to satisfy vocab size
    for id, _ in enumerate(range(vocab_size)):
        dummy_dict.add_symbol("{}".format(id), 1000)
    return dummy_dict


class DummyTask(LegacyFairseqTask):
    def __init__(self, args):
        super().__init__(args)
        self.dictionary = get_dummy_dictionary()
        if getattr(self.args, "ctc", False):
            self.dictionary.add_symbol("<ctc_blank>")
        self.tgt_dict = self.dictionary

    @property
    def target_dictionary(self):
        return self.dictionary


def get_dummy_task_and_parser():
    """
    to build a fariseq model, we need some dummy parse and task. This function
    is used to create dummy task and parser to faciliate model/criterion test

    Note: we use FbSpeechRecognitionTask as the dummy task. You may want
    to use other task by providing another function
    """
    parser = argparse.ArgumentParser(
        description="test_dummy_s2s_task", argument_default=argparse.SUPPRESS
    )
    DummyTask.add_args(parser)
    args = parser.parse_args([])
    task = DummyTask.setup_task(args)
    return task, parser


def get_dummy_input(T=100, D=80, B=5, K=100):
    forward_input = {}
    # T max sequence length
    # D feature vector dimension
    # B batch size
    # K target dimension size
    feature = torch.randn(B, T, D)
    # this (B, T, D) layout is just a convention, you can override it by
    # write your own _prepare_forward_input function
    src_lengths = torch.from_numpy(
        np.random.randint(low=1, high=T, size=B, dtype=np.int64)
    )
    src_lengths[0] = T  # make sure the maximum length matches
    prev_output_tokens = []
    for b in range(B):
        token_length = np.random.randint(low=1, high=src_lengths[b].item() + 1)
        tokens = np.random.randint(low=0, high=K, size=token_length, dtype=np.int64)
        prev_output_tokens.append(torch.from_numpy(tokens))

    prev_output_tokens = fairseq_data_utils.collate_tokens(
        prev_output_tokens,
        pad_idx=1,
        eos_idx=2,
        left_pad=False,
        move_eos_to_beginning=False,
    )
    src_lengths, sorted_order = src_lengths.sort(descending=True)
    forward_input["src_tokens"] = feature.index_select(0, sorted_order)
    forward_input["src_lengths"] = src_lengths
    forward_input["prev_output_tokens"] = prev_output_tokens

    return forward_input


def get_dummy_encoder_output(encoder_out_shape=(100, 80, 5)):
    """
    This only provides an example to generate dummy encoder output
    """
    (T, B, D) = encoder_out_shape
    encoder_out = {}

    encoder_out["encoder_out"] = torch.from_numpy(
        np.random.randn(*encoder_out_shape).astype(np.float32)
    )
    seq_lengths = torch.from_numpy(np.random.randint(low=1, high=T, size=B))
    # some dummy mask
    encoder_out["encoder_padding_mask"] = torch.arange(T).view(1, T).expand(
        B, -1
    ) >= seq_lengths.view(B, 1).expand(-1, T)
    encoder_out["encoder_padding_mask"].t_()

    # encoer_padding_mask is (T, B) tensor, with (t, b)-th element indicate
    # whether encoder_out[t, b] is valid (=0) or not (=1)
    return encoder_out


def _current_postion_info():
    cf = currentframe()
    frameinfo = " (at {}:{})".format(
        os.path.basename(getframeinfo(cf).filename), cf.f_back.f_lineno
    )
    return frameinfo


def check_encoder_output(encoder_output, batch_size=None):
    """we expect encoder_output to be a dict with the following
    key/value pairs:
    - encoder_out: a Torch.Tensor
    - encoder_padding_mask: a binary Torch.Tensor
    """
    if not isinstance(encoder_output, dict):
        msg = (
            "FairseqEncoderModel.forward(...) must be a dict" + _current_postion_info()
        )
        return False, msg

    if "encoder_out" not in encoder_output:
        msg = (
            "FairseqEncoderModel.forward(...) must contain encoder_out"
            + _current_postion_info()
        )
        return False, msg

    if "encoder_padding_mask" not in encoder_output:
        msg = (
            "FairseqEncoderModel.forward(...) must contain encoder_padding_mask"
            + _current_postion_info()
        )
        return False, msg

    if not isinstance(encoder_output["encoder_out"], torch.Tensor):
        msg = "encoder_out must be a torch.Tensor" + _current_postion_info()
        return False, msg

    if encoder_output["encoder_out"].dtype != torch.float32:
        msg = "encoder_out must have float32 dtype" + _current_postion_info()
        return False, msg

    mask = encoder_output["encoder_padding_mask"]
    if mask is not None:
        if not isinstance(mask, torch.Tensor):
            msg = (
                "encoder_padding_mask must be a torch.Tensor" + _current_postion_info()
            )
            return False, msg
        if mask.dtype != torch.uint8 and (
            not hasattr(torch, "bool") or mask.dtype != torch.bool
        ):
            msg = (
                "encoder_padding_mask must have dtype of uint8"
                + _current_postion_info()
            )
            return False, msg

        if mask.dim() != 2:
            msg = (
                "we expect encoder_padding_mask to be a 2-d tensor, in shape (T, B)"
                + _current_postion_info()
            )
            return False, msg

        if batch_size is not None and mask.size(1) != batch_size:
            msg = (
                "we expect encoder_padding_mask to be a 2-d tensor, with size(1)"
                + " being the batch size"
                + _current_postion_info()
            )
            return False, msg
    return True, None


def check_decoder_output(decoder_output):
    """we expect output from a decoder is a tuple with the following constraint:
    - the first element is a torch.Tensor
    - the second element can be anything (reserved for future use)
    """
    if not isinstance(decoder_output, tuple):
        msg = "FariseqDecoder output must be a tuple" + _current_postion_info()
        return False, msg

    if len(decoder_output) != 2:
        msg = "FairseqDecoder output must be 2-elem tuple" + _current_postion_info()
        return False, msg

    if not isinstance(decoder_output[0], torch.Tensor):
        msg = (
            "FariseqDecoder output[0] must be a torch.Tensor" + _current_postion_info()
        )
        return False, msg

    return True, None


# ///////////////////////////////////////////////////////////////////////////
# Base Test class
# ///////////////////////////////////////////////////////////////////////////


class TestBaseFairseqModelBase(unittest.TestCase):
    """
    This class is used to facilitate writing unittest for any class derived from
    `BaseFairseqModel`.
    """

    @classmethod
    def setUpClass(cls):
        if cls is TestBaseFairseqModelBase:
            raise unittest.SkipTest("Skipping test case in base")
        super().setUpClass()

    def setUpModel(self, model):
        self.assertTrue(isinstance(model, BaseFairseqModel))
        self.model = model

    def setupInput(self):
        pass

    def setUp(self):
        self.model = None
        self.forward_input = None
        pass


class TestFairseqEncoderDecoderModelBase(TestBaseFairseqModelBase):
    """
    base code to test FairseqEncoderDecoderModel (formally known as
    `FairseqModel`) must be derived from this base class
    """

    @classmethod
    def setUpClass(cls):
        if cls is TestFairseqEncoderDecoderModelBase:
            raise unittest.SkipTest("Skipping test case in base")
        super().setUpClass()

    def setUpModel(self, model_cls, extra_args_setters=None):
        self.assertTrue(
            issubclass(model_cls, (FairseqEncoderDecoderModel, FairseqModel)),
            msg="This class only tests for FairseqModel subclasses",
        )

        task, parser = get_dummy_task_and_parser()
        model_cls.add_args(parser)

        args = parser.parse_args([])

        if extra_args_setters is not None:
            for args_setter in extra_args_setters:
                args_setter(args)
        model = model_cls.build_model(args, task)
        self.model = model

    def setUpInput(self, input=None):
        self.forward_input = get_dummy_input() if input is None else input

    def setUp(self):
        super().setUp()

    def test_forward(self):
        if self.model and self.forward_input:
            forward_output = self.model.forward(**self.forward_input)
            # for FairseqEncoderDecoderModel, forward returns a tuple of two
            # elements, the first one is a Torch.Tensor
            succ, msg = check_decoder_output(forward_output)
            if not succ:
                self.assertTrue(succ, msg=msg)
            self.forward_output = forward_output

    def test_get_normalized_probs(self):
        if self.model and self.forward_input:
            forward_output = self.model.forward(**self.forward_input)
            logprob = self.model.get_normalized_probs(forward_output, log_probs=True)
            prob = self.model.get_normalized_probs(forward_output, log_probs=False)

            # in order for different models/criterion to play with each other
            # we need to know whether the logprob or prob output is batch_first
            # or not. We assume an additional attribute will be attached to logprob
            # or prob. If you find your code failed here, simply override
            # FairseqModel.get_normalized_probs, see example at
            # https://fburl.com/batch_first_example
            self.assertTrue(hasattr(logprob, "batch_first"))
            self.assertTrue(hasattr(prob, "batch_first"))

            self.assertTrue(torch.is_tensor(logprob))
            self.assertTrue(torch.is_tensor(prob))


class TestFairseqEncoderModelBase(TestBaseFairseqModelBase):
    """
    base class to test FairseqEncoderModel
    """

    @classmethod
    def setUpClass(cls):
        if cls is TestFairseqEncoderModelBase:
            raise unittest.SkipTest("Skipping test case in base")
        super().setUpClass()

    def setUpModel(self, model_cls, extra_args_setters=None):
        self.assertTrue(
            issubclass(model_cls, FairseqEncoderModel),
            msg="This class is only used for testing FairseqEncoderModel",
        )
        task, parser = get_dummy_task_and_parser()
        model_cls.add_args(parser)
        args = parser.parse_args([])
        if extra_args_setters is not None:
            for args_setter in extra_args_setters:
                args_setter(args)

        model = model_cls.build_model(args, task)
        self.model = model

    def setUpInput(self, input=None):
        self.forward_input = get_dummy_input() if input is None else input
        # get_dummy_input() is originally for s2s, here we delete extra dict
        # items, so it can be used for EncoderModel / Encoder as well
        self.forward_input.pop("prev_output_tokens", None)

    def setUp(self):
        super().setUp()

    def test_forward(self):
        if self.forward_input and self.model:
            bsz = self.forward_input["src_tokens"].size(0)
            forward_output = self.model.forward(**self.forward_input)

            # we expect forward_output to be a dict with the following
            # key/value pairs:
            # - encoder_out: a Torch.Tensor
            # - encoder_padding_mask: a binary Torch.Tensor
            succ, msg = check_encoder_output(forward_output, batch_size=bsz)
            if not succ:
                self.assertTrue(succ, msg=msg)
            self.forward_output = forward_output

    def test_get_normalized_probs(self):
        if self.model and self.forward_input:
            forward_output = self.model.forward(**self.forward_input)
            logprob = self.model.get_normalized_probs(forward_output, log_probs=True)
            prob = self.model.get_normalized_probs(forward_output, log_probs=False)

            # in order for different models/criterion to play with each other
            # we need to know whether the logprob or prob output is batch_first
            # or not. We assume an additional attribute will be attached to logprob
            # or prob. If you find your code failed here, simply override
            # FairseqModel.get_normalized_probs, see example at
            # https://fburl.com/batch_first_example
            self.assertTrue(hasattr(logprob, "batch_first"))
            self.assertTrue(hasattr(prob, "batch_first"))

            self.assertTrue(torch.is_tensor(logprob))
            self.assertTrue(torch.is_tensor(prob))


class TestFairseqEncoderBase(unittest.TestCase):
    """
    base class to test FairseqEncoder
    """

    @classmethod
    def setUpClass(cls):
        if cls is TestFairseqEncoderBase:
            raise unittest.SkipTest("Skipping test case in base")
        super().setUpClass()

    def setUpEncoder(self, encoder):
        self.assertTrue(
            isinstance(encoder, FairseqEncoder),
            msg="This class is only used for test FairseqEncoder",
        )
        self.encoder = encoder

    def setUpInput(self, input=None):
        self.forward_input = get_dummy_input() if input is None else input
        # get_dummy_input() is originally for s2s, here we delete extra dict
        # items, so it can be used for EncoderModel / Encoder as well
        self.forward_input.pop("prev_output_tokens", None)

    def setUp(self):
        self.encoder = None
        self.forward_input = None

    def test_forward(self):
        if self.encoder and self.forward_input:
            bsz = self.forward_input["src_tokens"].size(0)

            forward_output = self.encoder.forward(**self.forward_input)
            succ, msg = check_encoder_output(forward_output, batch_size=bsz)
            if not succ:
                self.assertTrue(succ, msg=msg)
            self.forward_output = forward_output


class TestFairseqDecoderBase(unittest.TestCase):
    """
    base class to test FairseqDecoder
    """

    @classmethod
    def setUpClass(cls):
        if cls is TestFairseqDecoderBase:
            raise unittest.SkipTest("Skipping test case in base")
        super().setUpClass()

    def setUpDecoder(self, decoder):
        self.assertTrue(
            isinstance(decoder, FairseqDecoder),
            msg="This class is only used for test FairseqDecoder",
        )
        self.decoder = decoder

    def setUpInput(self, input=None):
        self.forward_input = get_dummy_encoder_output() if input is None else input

    def setUpPrevOutputTokens(self, tokens=None):
        if tokens is None:
            self.encoder_input = get_dummy_input()
            self.prev_output_tokens = self.encoder_input["prev_output_tokens"]
        else:
            self.prev_output_tokens = tokens

    def setUp(self):
        self.decoder = None
        self.forward_input = None
        self.prev_output_tokens = None

    def test_forward(self):
        if (
            self.decoder is not None
            and self.forward_input is not None
            and self.prev_output_tokens is not None
        ):
            forward_output = self.decoder.forward(
                prev_output_tokens=self.prev_output_tokens,
                encoder_out=self.forward_input,
            )
            succ, msg = check_decoder_output(forward_output)
            if not succ:
                self.assertTrue(succ, msg=msg)
            self.forward_input = forward_output


class DummyEncoderModel(FairseqEncoderModel):
    def __init__(self, encoder):
        super().__init__(encoder)

    @classmethod
    def build_model(cls, args, task):
        return cls(DummyEncoder())

    def get_logits(self, net_output):
        # Inverse of sigmoid to use with BinaryCrossEntropyWithLogitsCriterion as
        # F.binary_cross_entropy_with_logits combines sigmoid and CE
        return torch.log(
            torch.div(net_output["encoder_out"], 1 - net_output["encoder_out"])
        )

    def get_normalized_probs(self, net_output, log_probs, sample=None):
        lprobs = super().get_normalized_probs(net_output, log_probs, sample=sample)
        lprobs.batch_first = True
        return lprobs


class DummyEncoder(FairseqEncoder):
    def __init__(self):
        super().__init__(None)

    def forward(self, src_tokens, src_lengths):
        mask, max_len = lengths_to_encoder_padding_mask(src_lengths)
        return {"encoder_out": src_tokens, "encoder_padding_mask": mask}


class CrossEntropyCriterionTestBase(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        if cls is CrossEntropyCriterionTestBase:
            raise unittest.SkipTest("Skipping base class test case")
        super().setUpClass()

    def setUpArgs(self):
        args = argparse.Namespace()
        args.sentence_avg = False
        args.threshold = 0.1  # to use with BinaryCrossEntropyWithLogitsCriterion
        return args

    def setUp(self):
        args = self.setUpArgs()
        self.model = DummyEncoderModel(encoder=DummyEncoder())
        self.criterion = self.criterion_cls.build_criterion(args, task=DummyTask(args))

    def get_src_tokens(self, correct_prediction, aggregate):
        """
        correct_prediction: True if the net_output (src_tokens) should
        predict the correct target
        aggregate: True if the criterion expects net_output (src_tokens)
        aggregated across time axis
        """
        predicted_idx = 0 if correct_prediction else 1
        if aggregate:
            src_tokens = torch.zeros((2, 2), dtype=torch.float)
            for b in range(2):
                src_tokens[b][predicted_idx] = 1.0
        else:
            src_tokens = torch.zeros((2, 10, 2), dtype=torch.float)
            for b in range(2):
                for t in range(10):
                    src_tokens[b][t][predicted_idx] = 1.0
        return src_tokens

    def get_target(self, soft_target):
        if soft_target:
            target = torch.zeros((2, 2), dtype=torch.float)
            for b in range(2):
                target[b][0] = 1.0
        else:
            target = torch.zeros((2, 10), dtype=torch.long)
        return target

    def get_test_sample(self, correct, soft_target, aggregate):
        src_tokens = self.get_src_tokens(correct, aggregate)
        target = self.get_target(soft_target)
        L = src_tokens.size(1)
        return {
            "net_input": {"src_tokens": src_tokens, "src_lengths": torch.tensor([L])},
            "target": target,
            "ntokens": src_tokens.size(0) * src_tokens.size(1),
        }