from typing import List

import torch
import torch.nn as nn


class Vocab:
    def __init__(self, char2idx=None, idx2char=None):
        if char2idx is None:
            char2idx = {
                "<pad>": 0,
                "<sos>": 1,
                "<eos>": 2,
                "<unk>": 3,
            }

        self.char2idx = {
            char: int(index)
            for char, index in char2idx.items()
        }

        if idx2char is None:
            self.idx2char = {
                index: char
                for char, index in self.char2idx.items()
            }
        else:
            self.idx2char = {
                int(index): char
                for index, char in idx2char.items()
            }

    def encode(self, text: str) -> List[int]:
        unk_id = self.char2idx["<unk>"]

        return [
            self.char2idx.get(character, unk_id)
            for character in text
        ]

    def decode(self, ids: List[int]) -> str:
        characters = []
        eos_id = self.char2idx["<eos>"]

        for index in ids:
            index = int(index)

            if index == eos_id:
                break

            if index > eos_id:
                characters.append(
                    self.idx2char.get(index, "")
                )

        return "".join(characters)


class LemmaModel(nn.Module):
    def __init__(
        self,
        vocab_size,
        char_emb_dim=96,
        hidden_size=128,
        drop_prob=0.30,
        num_heads=16,
        max_gen_len=30,
    ):
        super().__init__()

        self.max_gen_len = max_gen_len

        self.emb = nn.Embedding(
            vocab_size,
            char_emb_dim,
            padding_idx=0,
        )

        self.dropout_enc = nn.Dropout(drop_prob)
        self.dropout_dec = nn.Dropout(drop_prob)
        self.dropout_att = nn.Dropout(drop_prob)

        self.enc1 = nn.LSTM(
            char_emb_dim,
            hidden_size,
            bidirectional=True,
            batch_first=True,
        )

        self.enc2 = nn.LSTM(
            hidden_size * 2,
            hidden_size,
            bidirectional=True,
            batch_first=True,
        )

        self.attn = nn.MultiheadAttention(
            hidden_size * 2,
            num_heads,
            batch_first=True,
        )

        self.dec = nn.LSTM(
            char_emb_dim + hidden_size * 4,
            hidden_size * 2,
            batch_first=True,
        )

        self.dec_cross_attn = nn.MultiheadAttention(
            embed_dim=hidden_size * 2,
            num_heads=num_heads,
            kdim=hidden_size * 4,
            vdim=hidden_size * 4,
            batch_first=True,
        )

        self.out = nn.Linear(
            hidden_size * 2,
            vocab_size,
            bias=True,
        )

    def encode(self, src, src_lens):
        embedded = self.emb(src)

        packed1 = nn.utils.rnn.pack_padded_sequence(
            embedded,
            src_lens.cpu(),
            batch_first=True,
            enforce_sorted=False,
        )

        enc1_output, _ = self.enc1(packed1)

        enc1_output, _ = (
            nn.utils.rnn.pad_packed_sequence(
                enc1_output,
                batch_first=True,
            )
        )

        enc1_output = self.dropout_enc(enc1_output)

        packed2 = nn.utils.rnn.pack_padded_sequence(
            enc1_output,
            src_lens.cpu(),
            batch_first=True,
            enforce_sorted=False,
        )

        enc2_output, _ = self.enc2(packed2)

        enc2_output, _ = (
            nn.utils.rnn.pad_packed_sequence(
                enc2_output,
                batch_first=True,
            )
        )

        enc2_output = self.dropout_enc(enc2_output)

        attention_output, _ = self.attn(
            enc1_output,
            enc2_output,
            enc2_output,
        )

        attention_output = self.dropout_att(
            attention_output
        )

        return torch.cat(
            [enc2_output, attention_output],
            dim=-1,
        )

    def forward(self, src, src_lens, tgt):
        encoder_combined = self.encode(
            src,
            src_lens,
        )

        decoder_target = self.emb(tgt[:, :-1])
        target_len = decoder_target.size(1)

        if encoder_combined.size(1) >= target_len:
            combined_trimmed = encoder_combined[
                :, :target_len, :
            ]
        else:
            padding = encoder_combined.new_zeros(
                encoder_combined.size(0),
                target_len - encoder_combined.size(1),
                encoder_combined.size(2),
            )

            combined_trimmed = torch.cat(
                [encoder_combined, padding],
                dim=1,
            )

        decoder_input = torch.cat(
            [decoder_target, combined_trimmed],
            dim=-1,
        )

        decoder_output, _ = self.dec(decoder_input)
        decoder_output = self.dropout_dec(decoder_output)

        cross_output, _ = self.dec_cross_attn(
            decoder_output,
            encoder_combined,
            encoder_combined,
        )

        cross_output = self.dropout_att(cross_output)

        return self.out(cross_output)

    def generate(
        self,
        src,
        src_lens,
        vocab,
        max_len=None,
    ):
        self.eval()

        if max_len is None:
            max_len = self.max_gen_len

        batch_size = src.size(0)
        eos_id = vocab.char2idx["<eos>"]

        with torch.inference_mode():
            encoder_combined = self.encode(
                src,
                src_lens,
            )

            source_len = encoder_combined.size(1)

            current = torch.full(
                (batch_size, 1),
                vocab.char2idx["<sos>"],
                device=src.device,
                dtype=torch.long,
            )

            hidden = None
            hypotheses = [
                []
                for _ in range(batch_size)
            ]

            finished = torch.zeros(
                batch_size,
                dtype=torch.bool,
                device=src.device,
            )

            for step in range(max_len):
                embedded = self.emb(current).squeeze(1)

                combined_step = encoder_combined[
                    :,
                    min(step, source_len - 1),
                    :,
                ]

                decoder_input = torch.cat(
                    [embedded, combined_step],
                    dim=-1,
                ).unsqueeze(1)

                decoder_output, hidden = self.dec(
                    decoder_input,
                    hidden,
                )

                decoder_output = self.dropout_dec(
                    decoder_output
                )

                cross_output, _ = self.dec_cross_attn(
                    decoder_output,
                    encoder_combined,
                    encoder_combined,
                )

                cross_output = self.dropout_att(
                    cross_output
                )

                logits = self.out(cross_output)
                next_ids = logits.argmax(dim=-1)
                current = next_ids

                for index in range(batch_size):
                    if not finished[index]:
                        token_id = int(
                            next_ids[index, 0].item()
                        )

                        hypotheses[index].append(token_id)

                        if token_id == eos_id:
                            finished[index] = True

                if finished.all():
                    break

                current = torch.where(
                    finished.unsqueeze(1),
                    torch.full_like(current, eos_id),
                    current,
                )

        return [
            vocab.decode(hypothesis)
            for hypothesis in hypotheses
        ]