# Base Dependencies

# -----------------

from typing import Dict

# Package Dependencies

# --------------------

from .embeddings import Embeddings, RDEmbeddings

from .encoders import RecurrentEncoder

from .config import EmbeddingConfig, RDEmbeddingConfig, LSTMConfig

# Local Dependencies

# ------------------

from vocabulary import Vocabulary

# PyTorch Dependencies

# ---------------------

from torch import nn, Tensor, concat, mean

# Model

# -----

class HasanModel(nn.Module):

    """

    Implementation of the BiLSTM model described in `Hasan et al. (2020) - Integrating

    Text Embedding with Traditional NLP Features for Clinical Relation Extraction`

    """

    def __init__(

        self,

        vocab: Vocabulary,

        lstm_config: LSTMConfig,

        bioword2vec_config: EmbeddingConfig,

        rd_config: RDEmbeddingConfig,

        pos_config: EmbeddingConfig,

        dep_config: EmbeddingConfig,

        iob_config: EmbeddingConfig,

        num_classes: int = 2,

        clf_dropout: float = 0.25,

    ):

        """Initializes the model

        Args:

            vocab (Vocabulary): vocabulary object

            lstm_config (LSTMConfig): configuration for the LSTM encoder

            bioword2vec_config (EmbeddingConfig): configuration for the BioWord2Vec embedding

            rd_config (RDEmbeddingConfig): configuration for the Relative Distance embedding

            pos_config (EmbeddingConfig): configuration for the POS embedding

            dep_config (EmbeddingConfig): configuration for the DEP embedding

            iob_config (EmbeddingConfig): configuration for the IOB embedding

            num_classes (int, optional): number of output classes. Defaults to 2.

            clf_dropout (float, optional): dropout rate. Defaults to 0.1.

        """

        super(HasanModel, self).__init__()

        # attributes

        self.vocab = vocab

        self.lstm_config = lstm_config

        self.bioword2vec_config = bioword2vec_config

        self.rd_config = rd_config

        self.pos_config = pos_config

        self.dep_config = dep_config

        self.iob_config = iob_config

        self.num_classes = num_classes

        self.num_directions = 2 if self.lstm_config.bidirectional else 1

        self.clf_hidden_dim = 64

        self.clf_dropout = clf_dropout

        # embedding layers

        self.wv_embedding = Embeddings(**self.bioword2vec_config.__dict__)

        self.rd_embedding = RDEmbeddings(**self.rd_config.__dict__)

        self.pos_embedding = Embeddings(**self.pos_config.__dict__)

        self.dep_embedding = Embeddings(**self.dep_config.__dict__)

        self.iob_embedding = Embeddings(**self.iob_config.__dict__)

        # BiLSTM encoder

        self.lstm = RecurrentEncoder(rnn_type="lstm", **self.lstm_config.__dict__)

        # classifier

        self.fc = nn.Sequential(

            nn.Dropout(p=self.clf_dropout),

            nn.Linear(self.clf_input_dim, self.clf_hidden_dim),

            nn.ReLU(),

            nn.Dropout(p=self.clf_dropout),

            nn.Linear(self.clf_hidden_dim, self.num_classes),

            nn.ReLU(),

            nn.Sigmoid(),

        )

        # load pretrained embeddings

        self.wv_embedding.load_from_file(self.bioword2vec_config.emb_path, self.vocab)

    @property

    def clf_input_dim(self) -> int:

        """Input dimensions of the classifier"""

        return (self.num_directions * self.lstm_config.hidden_size) + (

            2 * self.wv_embedding.embedding_dim

        )

    def forward(self, inputs: Dict[str, Tensor]) -> Tensor:

        """Forward pass of the model

        Args:

            inputs (Dict[str, Tensor]): input tensors

        Returns:

            Tensor: output tensor

        """

        e1: Tensor = inputs["e1"]  # [batch_size, max_len_e1]

        e2: Tensor = inputs["e2"]  # [batch_size, max_len_e2]

        sent: Tensor = inputs["sent"]  # [batch_size, max_len_seq]

        rd1: Tensor = inputs["rd1"]  # [batch_size, max_len_seq]

        rd2: Tensor = inputs["rd2"]  # [batch_size, max_len_seq]

        pos: Tensor = inputs["pos"]  # [batch_size, max_len_seq]

        dep: Tensor = inputs["dep"]  # [batch_size, max_len_seq]

        iob: Tensor = inputs["iob"]  # [batch_size, max_len_seq]

        seq_length: Tensor = inputs["seq_length"]  # [batch_size]

        assert len(e1.shape) == 2

        assert len(e2.shape) == 2

        assert len(sent.shape) == 2

        assert len(rd1.shape) == 2

        assert len(rd2.shape) == 2

        assert len(pos.shape) == 2

        assert len(dep.shape) == 2 

        assert len(iob.shape) == 2

        assert len(seq_length.shape) == 1

        # embedded inputs

        e1_emb = mean(self.wv_embedding(e1), axis=1)  # [batch_size, wv_emb_dim]

        e2_emb = mean(self.wv_embedding(e2), axis=1)  # [batch_size, wv_emb_dim]

        sent_emb = self.wv_embedding(sent)  # [batch_size, seq_length, wv_emb_dim]

        rd1_emb = self.rd_embedding(rd1)  # [batch_size, seq_length, rd_emb_dim]

        rd2_emb = self.rd_embedding(rd2)  # [batch_size, seq_length, rd_emb_dim]

        pos_emb = self.pos_embedding(pos)  # [batch_size, seq_length, pos_emb_dim]

        dep_emb = self.dep_embedding(dep)  # [batch_size, seq_length, pos_emb_dim]

        iob_emb = self.iob_embedding(iob)  # [batch_size, seq_length, iob_emb_dim]

        # encode

        inputs_emb = concat((sent_emb, rd1_emb, rd2_emb, pos_emb, dep_emb, iob_emb), axis=2)

        outputs_emb, hidden_concat = self.lstm(inputs_emb, seq_length)

        outputs = concat((e1_emb, e2_emb, hidden_concat), axis=1)

        # classify

        outputs = self.fc(outputs)

        return outputs