import torch
import torch.nn as nn

class biLSTM(nn.Module):
    """
    The LSTM model that will be used to perform Sentiment analysis.
    """
    
    def __init__(self, 
                # объем словаря, с которым мы работаем, размер входа для слоя Embedding
                vocab_size: int, 
                # размер выходного эмбеддинга каждый элемент последовательности
                # будет описан вектором такой размерности
                embedding_dim: int, 
                # размерность hidden state LSTM слоя
                hidden_dim: int,
                # число слоев в LSTM
                n_layers: int,                
                drop_prob=0.5,
                seq_len = 128) -> None:
        
        super().__init__()
        self.hidden_dim = hidden_dim
        self.n_layers = n_layers
        self.seq_len = seq_len
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.lstm = nn.LSTM(embedding_dim, 
                            hidden_dim, 
                            n_layers, 
                            dropout=drop_prob, 
                            batch_first=True,
                            bidirectional=True
                            )
        
        self.do = nn.Dropout()
        
        self.fc1 = nn.Linear(2*hidden_dim * self.seq_len, 256)
        self.fc2 = nn.Linear(256, 1)
        self.sigmoid = nn.Sigmoid()
        
    def forward(self, x):
        embeds = self.embedding(x)
        lstm_out, _ = self.lstm(embeds)  
        out = self.fc2(torch.tanh(self.do(self.fc1(lstm_out.flatten(1)))))
        sig_out = self.sigmoid(out)
        
        return sig_out