import numpy as np
import torch
from torch import nn

from models.preprocess_stage.preprocess_lstm import preprocess_lstm

EMBEDDING_DIM = 128
HIDDEN_SIZE = 16
MAX_LEN = 125
# DEVICE='cpu'

embedding_matrix = np.load('models/datasets/embedding_matrix.npy')
embedding_layer = nn.Embedding.from_pretrained(torch.FloatTensor(embedding_matrix))


class AtenttionTest(nn.Module):
    def __init__(self, hidden_size=HIDDEN_SIZE):
        super().__init__()

        self.hidden_size = hidden_size
        self.fc1 = nn.Linear(self.hidden_size, self.hidden_size)
        self.fc2 = nn.Linear(self.hidden_size, self.hidden_size)
        self.tahn = nn.Tanh()
        self.fc3 = nn.Linear(self.hidden_size, 1)

    def forward(self, outputs_lmst, h_n):

        output_fc1 = self.fc1(outputs_lmst)
        output_fc2 = self.fc2(h_n.squeeze(0))

        fc1_fc2_cat = output_fc1 + output_fc2.unsqueeze(1)

        output_tahn = self.tahn(fc1_fc2_cat)

        attention_weights = torch.softmax(self.fc3(output_tahn).squeeze(2), dim=1)

        output_finished = torch.bmm(output_fc1.transpose(1, 2), attention_weights.unsqueeze(2))

        return output_finished, attention_weights


class LSTMnn(nn.Module):

    def __init__(self):
        super().__init__()

        self.embedding = embedding_layer
        self.lstm = nn.LSTM(
            input_size=EMBEDDING_DIM,
            hidden_size=HIDDEN_SIZE,
            num_layers=1,
            batch_first=True
        )
        self.attention = AtenttionTest(hidden_size=HIDDEN_SIZE)
        self.fc_out = nn.Sequential(
            nn.Linear(HIDDEN_SIZE, 128),
            nn.Dropout(),
            nn.Tanh(),
            nn.Linear(128, 1)
        )

    def forward(self, x):

        embedding = self.embedding(x)

        output_lstm, (h_n, _) = self.lstm(embedding)

        output_attention, attention_weights = self.attention(output_lstm, h_n)

        output_finished = self.fc_out(output_attention.squeeze(2))

        return torch.sigmoid(output_finished), attention_weights


model = LSTMnn()
model.load_state_dict(torch.load('models/weights/LSTMBestWeights.pt', map_location=torch.device('cpu')))


def predict_3(text):

    preprocessed_text = preprocess_lstm(text, MAX_LEN=MAX_LEN)

    # model.to(DEVICE)
    model.eval()
    predict, attention = model(torch.tensor(preprocessed_text).unsqueeze(0))

    predict = round(predict.item())

    return predict