#!/usr/bin/env python3

import io
import logging
import unittest

import numpy as np
import torch
from pytorch_translate import char_source_hybrid  # noqa
from pytorch_translate import char_source_model  # noqa
from pytorch_translate import char_source_transformer_model  # noqa
from pytorch_translate import rnn  # noqa
from pytorch_translate import transformer  # noqa
from pytorch_translate import constants
from pytorch_translate.ensemble_export import BeamSearch, BeamSearchAndDecode
from pytorch_translate.research.knowledge_distillation import (  # noqa
    dual_decoder_kd_model,
    hybrid_dual_decoder_kd_model,
)
from pytorch_translate.tasks import pytorch_translate_task as tasks
from pytorch_translate.test import utils as test_utils


logger = logging.getLogger(__name__)


class TestPyTorchExport(unittest.TestCase):
    def _test_full_beam_decoder(self, test_args, quantize=False):
        samples, src_dict, tgt_dict = test_utils.prepare_inputs(test_args)
        task = tasks.DictionaryHolderTask(src_dict, tgt_dict)
        sample = next(samples)
        # [seq len, batch size=1]
        src_tokens = sample["net_input"]["src_tokens"][0:1].t()
        # [seq len]
        src_lengths = sample["net_input"]["src_lengths"][0:1].long()

        num_models = 3
        model_list = []
        for _ in range(num_models):
            model_list.append(task.build_model(test_args))

        length, word_length = 11, 7
        if test_args.arch in constants.ARCHS_FOR_CHAR_SOURCE:
            char_inds = torch.LongTensor(
                np.random.randint(0, 126, (1, length, word_length), dtype="int64")
            )
            word_lengths = torch.IntTensor(
                np.array([word_length] * length, dtype="int32")
            ).reshape((1, length))
        else:
            char_inds, word_lengths = None, None

        beam_size = 6
        bs = BeamSearch(
            model_list,
            tgt_dict,
            src_tokens,
            src_lengths,
            beam_size=beam_size,
            quantize=quantize,
            char_inds=char_inds,
            word_lengths=word_lengths,
        )
        f = io.BytesIO()
        bs.save_to_pytorch(f)

        # Test generalization with a different sequence length
        src_tokens = torch.LongTensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]).unsqueeze(1)
        src_lengths = torch.LongTensor([11])
        prev_token = torch.LongTensor([0])
        prev_scores = torch.FloatTensor([0.0])
        attn_weights = torch.zeros(src_tokens.shape[0])
        prev_hypos_indices = torch.zeros(beam_size, dtype=torch.int64)

        outs = bs(
            src_tokens,
            src_lengths,
            prev_token,
            prev_scores,
            attn_weights,
            prev_hypos_indices,
            torch.LongTensor([20]),
            char_inds=char_inds,
            word_lengths=word_lengths,
        )

        f.seek(0)
        deserialized_bs = torch.jit.load(f)
        deserialized_bs.apply(lambda s: s._unpack() if hasattr(s, "_unpack") else None)
        outs_deserialized = deserialized_bs(
            src_tokens,
            src_lengths,
            prev_token,
            prev_scores,
            attn_weights,
            prev_hypos_indices,
            torch.LongTensor([20]),
            char_inds=char_inds,
            word_lengths=word_lengths,
        )

        for a, b in zip(outs_deserialized, outs):
            np.testing.assert_allclose(a.detach().numpy(), b.detach().numpy())

    def test_full_beam_decoder(self):
        test_args = test_utils.ModelParamsDict(
            encoder_bidirectional=True, sequence_lstm=True
        )
        self._test_full_beam_decoder(test_args)

    def test_full_beam_decoder_vocab_reduction(self):
        test_args = test_utils.ModelParamsDict(
            encoder_bidirectional=True, sequence_lstm=True
        )
        lexical_dictionaries = test_utils.create_lexical_dictionaries()
        test_args.vocab_reduction_params = {
            "lexical_dictionaries": lexical_dictionaries,
            "num_top_words": 10,
            "max_translation_candidates_per_word": 1,
        }
        self._test_full_beam_decoder(test_args)

    def test_full_beam_decoder_quantization(self):
        test_args = test_utils.ModelParamsDict(
            encoder_bidirectional=True, sequence_lstm=True
        )
        lexical_dictionaries = test_utils.create_lexical_dictionaries()
        test_args.vocab_reduction_params = {
            "lexical_dictionaries": lexical_dictionaries,
            "num_top_words": 10,
            "max_translation_candidates_per_word": 1,
        }
        self._test_full_beam_decoder(test_args, quantize=True)

    def test_full_beam_decoder_char_rnn_vocab_reduction(self):
        test_args = test_utils.ModelParamsDict(
            encoder_bidirectional=True, sequence_lstm=True
        )
        lexical_dictionaries = test_utils.create_lexical_dictionaries()
        test_args.vocab_reduction_params = {
            "lexical_dictionaries": lexical_dictionaries,
            "num_top_words": 10,
            "max_translation_candidates_per_word": 1,
        }

        test_args.arch = "char_source"
        test_args.char_source_dict_size = 126
        test_args.char_embed_dim = 8
        test_args.char_rnn_units = 12
        test_args.char_rnn_layers = 2

        self._test_full_beam_decoder(test_args)

    def test_full_beam_decoder_char_cnn_vocab_reduction(self):
        test_args = test_utils.ModelParamsDict(
            encoder_bidirectional=True, sequence_lstm=True
        )
        lexical_dictionaries = test_utils.create_lexical_dictionaries()
        test_args.vocab_reduction_params = {
            "lexical_dictionaries": lexical_dictionaries,
            "num_top_words": 10,
            "max_translation_candidates_per_word": 1,
        }

        test_args.arch = "char_source"
        test_args.char_source_dict_size = 126
        test_args.char_embed_dim = 8
        test_args.char_cnn_params = "[(10, 3), (10, 5)]"
        test_args.char_cnn_nonlinear_fn = "tanh"
        test_args.char_cnn_pool_type = "max"
        test_args.char_cnn_num_highway_layers = 2

        self._test_full_beam_decoder(test_args)

    def test_full_beam_decoder_quantization_hybrid(self):
        test_args = test_utils.ModelParamsDict(arch="hybrid_transformer_rnn")
        lexical_dictionaries = test_utils.create_lexical_dictionaries()
        test_args.vocab_reduction_params = {
            "lexical_dictionaries": lexical_dictionaries,
            "num_top_words": 10,
            "max_translation_candidates_per_word": 1,
        }
        self._test_full_beam_decoder(test_args, quantize=True)

    def test_full_beam_decoder_quantization_hybrid_reduced_attention(self):
        test_args = test_utils.ModelParamsDict(arch="hybrid_transformer_rnn")
        test_args.decoder_reduced_attention_dim = 10
        lexical_dictionaries = test_utils.create_lexical_dictionaries()
        test_args.vocab_reduction_params = {
            "lexical_dictionaries": lexical_dictionaries,
            "num_top_words": 10,
            "max_translation_candidates_per_word": 1,
        }
        self._test_full_beam_decoder(test_args, quantize=True)

    def test_full_beam_decoder_aan(self):
        test_args = test_utils.ModelParamsDict(arch="transformer")
        test_args.aan = True
        lexical_dictionaries = test_utils.create_lexical_dictionaries()
        test_args.vocab_reduction_params = {
            "lexical_dictionaries": lexical_dictionaries,
            "num_top_words": 10,
            "max_translation_candidates_per_word": 1,
        }
        self._test_full_beam_decoder(test_args, quantize=True)

    def test_full_beam_decoder_aan_bottlenceck(self):
        test_args = test_utils.ModelParamsDict(arch="transformer")
        test_args.aan = True
        test_args.decoder_out_embed_dim = 5
        lexical_dictionaries = test_utils.create_lexical_dictionaries()
        test_args.vocab_reduction_params = {
            "lexical_dictionaries": lexical_dictionaries,
            "num_top_words": 10,
            "max_translation_candidates_per_word": 1,
        }
        self._test_full_beam_decoder(test_args, quantize=True)

    def test_full_beam_decoder_char_transformer(self):
        test_args = test_utils.ModelParamsDict(arch="char_transformer")
        test_args.char_source_dict_size = 126
        test_args.char_embed_dim = 8
        test_args.char_cnn_params = "[(10, 3), (10, 5)]"
        test_args.char_cnn_nonlinear_fn = "tanh"
        test_args.char_cnn_pool_type = "max"
        test_args.char_cnn_num_highway_layers = 2

        self._test_full_beam_decoder(test_args)

    def test_full_beam_decoder_char_source_hybrid(self):
        test_args = test_utils.ModelParamsDict(arch="char_source_hybrid")
        test_args.char_source_dict_size = 126
        test_args.char_embed_dim = 8
        test_args.char_cnn_params = "[(8, 3), (8, 5)]"
        test_args.char_cnn_nonlinear_fn = "tanh"
        test_args.char_cnn_pool_type = "max"
        test_args.char_cnn_num_highway_layers = 2
        self._test_full_beam_decoder(test_args)


class TestBeamSearchAndDecodeExport(unittest.TestCase):
    def _test_full_beam_search_decoder(self, test_args, quantize=False):
        samples, src_dict, tgt_dict = test_utils.prepare_inputs(test_args)
        task = tasks.DictionaryHolderTask(src_dict, tgt_dict)
        sample = next(samples)
        # [seq len, batch size=1]
        src_tokens = sample["net_input"]["src_tokens"][0:1].t()
        # [seq len]
        src_lengths = sample["net_input"]["src_lengths"][0:1].long()

        num_models = 3
        model_list = []
        for _ in range(num_models):
            model_list.append(task.build_model(test_args))

        eos_token_id = 8
        length_penalty = 0.25
        nbest = 3
        stop_at_eos = True
        num_steps = torch.LongTensor([20])

        beam_size = 6
        bsd = BeamSearchAndDecode(
            model_list,
            tgt_dict,
            src_tokens,
            src_lengths,
            eos_token_id=eos_token_id,
            length_penalty=length_penalty,
            nbest=nbest,
            beam_size=beam_size,
            stop_at_eos=stop_at_eos,
            quantize=quantize,
        )
        f = io.BytesIO()
        bsd.save_to_pytorch(f)

        # Test generalization with a different sequence length
        src_tokens = torch.LongTensor([1, 2, 3, 4, 5, 6, 7, 9, 9, 10, 11]).unsqueeze(1)
        src_lengths = torch.LongTensor([11])
        prev_token = torch.LongTensor([0])
        prev_scores = torch.FloatTensor([0.0])
        attn_weights = torch.zeros(src_tokens.shape[0])
        prev_hypos_indices = torch.zeros(beam_size, dtype=torch.int64)

        outs = bsd(
            src_tokens,
            src_lengths,
            prev_token,
            prev_scores,
            attn_weights,
            prev_hypos_indices,
            num_steps[0],
        )

        f.seek(0)
        deserialized_bsd = torch.jit.load(f)
        deserialized_bsd.apply(lambda s: s._unpack() if hasattr(s, "_unpack") else None)
        outs_deserialized = deserialized_bsd(
            src_tokens,
            src_lengths,
            prev_token,
            prev_scores,
            attn_weights,
            prev_hypos_indices,
            num_steps[0],
        )

        for hypo, hypo_deserialized in zip(outs, outs_deserialized):
            np.testing.assert_array_equal(
                hypo[0].tolist(), hypo_deserialized[0].tolist()
            )
            np.testing.assert_array_almost_equal(
                hypo[2], hypo_deserialized[2], decimal=1
            )
            np.testing.assert_array_almost_equal(
                hypo[3].numpy(), hypo_deserialized[3].numpy(), decimal=1
            )

    def test_full_beam_search_decoder(self):
        test_args = test_utils.ModelParamsDict(
            encoder_bidirectional=True, sequence_lstm=True
        )
        self._test_full_beam_search_decoder(test_args)

    def test_full_beam_search_decoder_reverse(self):
        test_args = test_utils.ModelParamsDict(
            encoder_bidirectional=True, sequence_lstm=True, reverse_source=True
        )
        self._test_full_beam_search_decoder(test_args)

    def test_full_beam_search_decoder_vocab_reduction(self):
        test_args = test_utils.ModelParamsDict(
            encoder_bidirectional=True, sequence_lstm=True
        )
        lexical_dictionaries = test_utils.create_lexical_dictionaries()
        test_args.vocab_reduction_params = {
            "lexical_dictionaries": lexical_dictionaries,
            "num_top_words": 10,
            "max_translation_candidates_per_word": 1,
        }
        self._test_full_beam_search_decoder(test_args)

    def test_full_beam_search_decoder_quantization(self):
        test_args = test_utils.ModelParamsDict(
            encoder_bidirectional=True, sequence_lstm=True
        )
        lexical_dictionaries = test_utils.create_lexical_dictionaries()
        test_args.vocab_reduction_params = {
            "lexical_dictionaries": lexical_dictionaries,
            "num_top_words": 10,
            "max_translation_candidates_per_word": 1,
        }
        self._test_full_beam_search_decoder(test_args, quantize=True)

    def test_full_beam_search_decoder_hybrid(self):
        test_args = test_utils.ModelParamsDict(arch="hybrid_transformer_rnn")
        self._test_full_beam_search_decoder(test_args)

    def test_full_beam_search_decoder_quantization_hybrid(self):
        test_args = test_utils.ModelParamsDict(arch="hybrid_transformer_rnn")
        lexical_dictionaries = test_utils.create_lexical_dictionaries()
        test_args.vocab_reduction_params = {
            "lexical_dictionaries": lexical_dictionaries,
            "num_top_words": 10,
            "max_translation_candidates_per_word": 1,
        }
        self._test_full_beam_search_decoder(test_args, quantize=True)
