import unittest
import torch
from nlp_core.vocabulary import Vocabulary
from models.abstractive.attention import BahdanauAttention, LuongAttention
from models.abstractive.encoder import Seq2SeqEncoder
from models.abstractive.decoder import Seq2SeqDecoder
from models.abstractive.seq2seq_model import Seq2SeqSummarizer

class TestAbstractiveModel(unittest.TestCase):

    def setUp(self):
        self.src_vocab = Vocabulary("src_vocab")
        self.src_vocab.add_sentence(["natural", "language", "processing", "is", "great"])
        self.src_vocab.build_vocab(min_freq=1)

        self.trg_vocab = Vocabulary("trg_vocab")
        self.trg_vocab.add_sentence(["nlp", "is", "great"])
        self.trg_vocab.build_vocab(min_freq=1)

        self.emb_dim = 32
        self.enc_hid_dim = 64
        self.dec_hid_dim = 64

    def test_attention_shapes(self):
        batch_size = 2
        src_len = 5
        hidden = torch.randn(batch_size, self.dec_hid_dim)
        encoder_outputs = torch.randn(batch_size, src_len, self.enc_hid_dim * 2)

        bahdanau = BahdanauAttention(self.enc_hid_dim * 2, self.dec_hid_dim)
        attn_weights = bahdanau(hidden, encoder_outputs)
        self.assertEqual(attn_weights.shape, (batch_size, src_len))
        self.assertAlmostEqual(attn_weights.sum(dim=1)[0].item(), 1.0, places=4)

    def test_seq2seq_forward_and_decoding(self):
        model = Seq2SeqSummarizer(
            src_vocab=self.src_vocab,
            trg_vocab=self.trg_vocab,
            emb_dim=self.emb_dim,
            enc_hid_dim=self.enc_hid_dim,
            dec_hid_dim=self.dec_hid_dim
        )

        batch_size = 2
        src_len = 6
        trg_len = 4
        src = torch.randint(0, len(self.src_vocab), (batch_size, src_len))
        trg = torch.randint(0, len(self.trg_vocab), (batch_size, trg_len))

        outputs = model(src, trg, teacher_forcing_ratio=0.5)
        self.assertEqual(outputs.shape, (batch_size, trg_len, len(self.trg_vocab)))

        # Test Greedy generation
        summary_tokens, attn = model.summarize_greedy(["natural", "language", "processing"], max_len=5)
        self.assertIsInstance(summary_tokens, list)

        # Test Beam Search generation
        beam_summary = model.summarize_beam(["natural", "language", "processing"], beam_width=2, max_len=5)
        self.assertIsInstance(beam_summary, list)

if __name__ == "__main__":
    unittest.main()