import unittest
from fastNLP.models import SequenceGeneratorModel
from fastNLP.models import LSTMSeq2SeqModel, TransformerSeq2SeqModel
from fastNLP import Vocabulary, DataSet
import torch
from fastNLP.embeddings import StaticEmbedding
from fastNLP import Trainer, CrossEntropyLoss, AccuracyMetric
from fastNLP import Callback


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)

    src_words_idx = [[3, 1, 2], [1, 2]]
    # tgt_words_idx = [[1, 2, 3, 4], [2, 3]]
    src_seq_len = [3, 2]
    # tgt_seq_len = [4, 2]

    ds = DataSet({'src_tokens': src_words_idx, 'src_seq_len': src_seq_len, 'tgt_tokens': src_words_idx,
                  'tgt_seq_len':src_seq_len})

    ds.set_input('src_tokens', 'tgt_tokens', 'src_seq_len')
    ds.set_target('tgt_seq_len', 'tgt_tokens')

    return embed, ds


class ExitCallback(Callback):
    def __init__(self):
        super().__init__()

    def on_valid_end(self, eval_result, metric_key, optimizer, is_better_eval):
        if eval_result['AccuracyMetric']['acc']==1:
            raise KeyboardInterrupt()


class TestSeq2SeqGeneratorModel(unittest.TestCase):
    def test_run(self):
        #  检测是否能够使用SequenceGeneratorModel训练, 透传预测
        embed, ds = prepare_env()
        model1 = TransformerSeq2SeqModel.build_model(src_embed=embed, tgt_embed=None,
                                                    pos_embed='sin', max_position=20, num_layers=2, d_model=30, n_head=6,
                                                    dim_ff=20, dropout=0.1,
                                                    bind_encoder_decoder_embed=True,
                                                    bind_decoder_input_output_embed=True)
        trainer = Trainer(ds, model1, optimizer=None, loss=CrossEntropyLoss(target='tgt_tokens', seq_len='tgt_seq_len'),
                 batch_size=32, sampler=None, drop_last=False, update_every=1,
                 num_workers=0, n_epochs=100, print_every=5,
                 dev_data=ds, metrics=AccuracyMetric(target='tgt_tokens', seq_len='tgt_seq_len'), metric_key=None,
                 validate_every=-1, save_path=None, use_tqdm=False, device=None,
                 callbacks=ExitCallback(), check_code_level=0)
        res = trainer.train()
        self.assertEqual(res['best_eval']['AccuracyMetric']['acc'], 1)

        embed, ds = prepare_env()
        model2 = LSTMSeq2SeqModel.build_model(src_embed=embed, tgt_embed=None,
                                              num_layers=1, hidden_size=20, dropout=0.1,
                                              bind_encoder_decoder_embed=True,
                                              bind_decoder_input_output_embed=True, attention=True)
        optimizer = torch.optim.Adam(model2.parameters(), lr=0.01)
        trainer = Trainer(ds, model2, optimizer=optimizer, loss=CrossEntropyLoss(target='tgt_tokens', seq_len='tgt_seq_len'),
                          batch_size=32, sampler=None, drop_last=False, update_every=1,
                          num_workers=0, n_epochs=200, print_every=1,
                          dev_data=ds, metrics=AccuracyMetric(target='tgt_tokens', seq_len='tgt_seq_len'),
                          metric_key=None,
                          validate_every=-1, save_path=None, use_tqdm=False, device=None,
                          callbacks=ExitCallback(), check_code_level=0)
        res = trainer.train()
        self.assertEqual(res['best_eval']['AccuracyMetric']['acc'], 1)