import unittest

import torch
from fastNLP.modules.generator import SequenceGenerator
from fastNLP.modules import TransformerSeq2SeqDecoder, LSTMSeq2SeqDecoder, Seq2SeqDecoder, State
from fastNLP import Vocabulary
from fastNLP.embeddings import StaticEmbedding
from torch import nn
from fastNLP import seq_len_to_mask


def prepare_env():
    vocab = Vocabulary().add_word_lst("This is a test .".split())
    vocab.add_word_lst("Another test !".split())
    embed = StaticEmbedding(vocab, model_dir_or_name=None, embedding_dim=5)

    encoder_output = torch.randn(2, 3, 10)
    src_seq_len = torch.LongTensor([3, 2])
    encoder_mask = seq_len_to_mask(src_seq_len)

    return embed, encoder_output, encoder_mask


class TestSequenceGenerator(unittest.TestCase):
    def test_run(self):
        # 测试能否运行 (1) 初始化decoder，(2) decode一发
        embed, encoder_output, encoder_mask = prepare_env()

        for do_sample in [True, False]:
            for num_beams in [1, 3, 5]:
                with self.subTest(do_sample=do_sample, num_beams=num_beams):
                    decoder = LSTMSeq2SeqDecoder(embed=embed, num_layers=1, hidden_size=10,
                             dropout=0.3, bind_decoder_input_output_embed=True, attention=True)
                    state = decoder.init_state(encoder_output, encoder_mask)
                    generator = SequenceGenerator(decoder=decoder, max_length=20, num_beams=num_beams,
                             do_sample=do_sample, temperature=1.0, top_k=50, top_p=1.0, bos_token_id=1, eos_token_id=None,
                             repetition_penalty=1, length_penalty=1.0, pad_token_id=0)
                    generator.generate(state=state, tokens=None)

                    decoder = TransformerSeq2SeqDecoder(embed=embed, pos_embed=nn.Embedding(10, embed.embedding_dim),
                             d_model=encoder_output.size(-1), num_layers=2, n_head=2, dim_ff=10, dropout=0.1,
                             bind_decoder_input_output_embed=True)
                    state = decoder.init_state(encoder_output, encoder_mask)
                    generator = SequenceGenerator(decoder=decoder, max_length=5, num_beams=num_beams,
                             do_sample=do_sample, temperature=1.0, top_k=50, top_p=1.0, bos_token_id=1, eos_token_id=None,
                             repetition_penalty=1, length_penalty=1.0, pad_token_id=0)
                    generator.generate(state=state, tokens=None)

                    # 测试一下其它值
                    decoder = TransformerSeq2SeqDecoder(embed=embed, pos_embed=nn.Embedding(10, embed.embedding_dim),
                                                        d_model=encoder_output.size(-1), num_layers=2, n_head=2, dim_ff=10,
                                                        dropout=0.1,
                                                        bind_decoder_input_output_embed=True)
                    state = decoder.init_state(encoder_output, encoder_mask)
                    generator = SequenceGenerator(decoder=decoder, max_length=5, num_beams=num_beams,
                                                  do_sample=do_sample, temperature=0.9, top_k=50, top_p=0.5, bos_token_id=1,
                                                  eos_token_id=3, repetition_penalty=2, length_penalty=1.5, pad_token_id=0)
                    generator.generate(state=state, tokens=None)

    def test_greedy_decode(self):
        # 测试能否正确的generate
        class GreedyDummyDecoder(Seq2SeqDecoder):
            def __init__(self, decoder_output):
                super().__init__()
                self.cur_length = 0
                self.decoder_output = decoder_output

            def decode(self, tokens, state):
                self.cur_length += 1
                scores = self.decoder_output[:, self.cur_length]
                return scores

        class DummyState(State):
            def __init__(self, decoder):
                super().__init__()
                self.decoder = decoder

            def reorder_state(self, indices: torch.LongTensor):
                self.decoder.decoder_output = self._reorder_state(self.decoder.decoder_output, indices, dim=0)

        # greedy
        for beam_search in [1, 3]:
            decoder_output = torch.randn(2, 10, 5)
            path = decoder_output.argmax(dim=-1)  # 2 x 4
            decoder = GreedyDummyDecoder(decoder_output)
            with self.subTest(beam_search=beam_search):
                generator = SequenceGenerator(decoder=decoder, max_length=decoder_output.size(1), num_beams=beam_search,
                                                      do_sample=False, temperature=1, top_k=50, top_p=1, bos_token_id=1,
                                                      eos_token_id=None, repetition_penalty=1, length_penalty=1, pad_token_id=0)
                decode_path = generator.generate(DummyState(decoder), tokens=decoder_output[:, 0].argmax(dim=-1, keepdim=True))

                self.assertEqual(decode_path.eq(path).sum(), path.numel())

        # greedy check eos_token_id
        for beam_search in [1, 3]:
            decoder_output = torch.randn(2, 10, 5)
            decoder_output[:, :7, 4].fill_(-100)
            decoder_output[0, 7, 4] = 1000  # 在第8个结束
            decoder_output[1, 5, 4] = 1000
            path = decoder_output.argmax(dim=-1)  # 2 x 4
            decoder = GreedyDummyDecoder(decoder_output)
            with self.subTest(beam_search=beam_search):
                generator = SequenceGenerator(decoder=decoder, max_length=decoder_output.size(1), num_beams=beam_search,
                                              do_sample=False, temperature=1, top_k=50, top_p=0.5, bos_token_id=1,
                                              eos_token_id=4, repetition_penalty=1, length_penalty=1, pad_token_id=0)
                decode_path = generator.generate(DummyState(decoder),
                                                 tokens=decoder_output[:, 0].argmax(dim=-1, keepdim=True))
                self.assertEqual(decode_path.size(1), 8)  # 长度为8
                self.assertEqual(decode_path[0].eq(path[0, :8]).sum(), 8)
                self.assertEqual(decode_path[1, :6].eq(path[1, :6]).sum(), 6)