import json
import numpy as np


import torch
import torch.nn.functional as F
import torch.nn as nn

import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings('ignore')

from gensim.models import Word2Vec
from typing import Tuple
from time import time
import string
from pymystem3 import Mystem
import pickle
import joblib

sw = ['х', 'эх', 'это','и', 'в', 'во', 'не', 'что', 'он', 'на', 'я', 'с', 'со', 'как', 'а', 'то', 'все', 'она', 'так', 'его', 'но', 'да', 'ты', 'к', 'у', 'же', 'вы', 'за', 'бы', 'по', 'только', 'ее', 'мне', 'было', 'вот', 'от', 'меня', 'еще', 'нет', 'о', 'из', 'ему', 'теперь', 'когда', 'даже', 'ну', 'вдруг', 'ли', 'если', 'уже', 'или', 'ни', 'быть', 'был', 'него', 'до', 'вас', 'нибудь', 'опять', 'уж', 'вам', 'ведь', 'там', 'потом', 'себя', 'ничего', 'ей', 'может', 'они', 'тут', 'где', 'есть', 'надо', 'ней', 'для', 'мы', 'тебя', 'их', 'чем', 'была', 'сам', 'чтоб', 'без', 'будто', 'чего', 'раз', 'тоже', 'себе', 'под', 'будет', 'ж', 'тогда', 'кто', 'этот', 'того', 'потому', 'этого', 'какой', 'совсем', 'ним', 'здесь', 'этом', 'один', 'почти', 'мой', 'тем', 'чтобы', 'нее', 'сейчас', 'были', 'куда', 'зачем', 'всех', 'никогда', 'можно', 'при', 'наконец', 'два', 'об', 'другой', 'хоть', 'после', 'над', 'больше', 'тот', 'через', 'эти', 'нас', 'про', 'всего', 'них', 'какая', 'много', 'разве', 'три', 'эту', 'моя', 'впрочем', 'хорошо', 'свою', 'этой', 'перед', 'иногда', 'лучше', 'чуть', 'том', 'нельзя', 'такой', 'им', 'более', 'всегда', 'конечно', 'всю', 'между']

punc = string.punctuation + '«»—…' + string.digits
mystem = Mystem()

def clean(text):
    
    text = text.lower() # нижний регистр
    text = "".join([c for c in text if c not in punc])
    text = mystem.lemmatize(text)
    text = ' '.join([word for word in text if word not in sw])
    text = " ".join(text.split()) # удаляем /n
    return text



with open('srcs/vocab_to_int.json', encoding='utf-8') as f:
    vocab_to_int = json.load(f)
VOCAB_SIZE = len(vocab_to_int) + 1
EMBEDDING_DIM = 64 # embedding_dim
SEQ_LEN = 350
HIDDEN_SIZE = 64


with open('srcs/embedding_matrix.npy', 'rb') as f:
    embedding_matrix = np.load(f)
embedding_layer = torch.nn.Embedding.from_pretrained(torch.FloatTensor(embedding_matrix))

class BahdanauAttention(nn.Module):
    def __init__(self, hidden_size: int) -> None:
        super().__init__()

        self.W_q = nn.Linear(hidden_size, hidden_size)
        self.W_k = nn.Linear(2*hidden_size, hidden_size)
        self.V = nn.Linear(hidden_size, 1)

    def forward(
            self,
            keys: torch.Tensor,
            query: torch.Tensor
            ) -> Tuple[torch.Tensor, torch.Tensor]:
        

        query = self.W_q(query)
        
        keys = self.W_k(keys)

        energy = self.V(F.tanh(query[0].unsqueeze(1) + query[1].unsqueeze(1) + keys)).squeeze(-1)
       
        weights = torch.softmax(energy, -1)
        context = torch.bmm(weights.unsqueeze(1), keys).squeeze(0)
        
        return context, weights    
    
class LSTMConcatAttentionB(nn.Module):

    def __init__(self) -> None:
        super().__init__()

        self.embedding = embedding_layer
        self.lstm = nn.LSTM(EMBEDDING_DIM, int(HIDDEN_SIZE), batch_first=True, bidirectional=True)
        self.attnB = BahdanauAttention(HIDDEN_SIZE)
        
        self.clf = nn.Sequential(
            nn.Linear(HIDDEN_SIZE, 512),
            nn.Dropout(),
            nn.Tanh(),
            nn.Linear(512, 256),
            nn.SELU(),
            nn.Linear(256, 3),
        )
    
    def forward(self, x):
        embeddings = self.embedding(x)
        
        outputs, (h_n, _) = self.lstm(embeddings)

        att_hidden, att_weights = self.attnB(outputs, h_n.squeeze(0))
        
        out = self.clf(att_hidden).squeeze()
        
        return out, att_weights

class Text_ex():
    def __init__(self, clean_func, voc, sl) -> None:
        self.clean = clean_func
        self.vocab = voc
        self.seq_len = sl

    def __call__(self, text):
        c_text = self.clean(text)
        review = [self.vocab[word] for word in c_text.split() if vocab_to_int.get(word)]
        if len(review) <= self.seq_len:
            zeros = list(np.zeros(self.seq_len - len(review)))
            new = zeros + review
        else:
            new = review[-self.seq_len:]
        t = torch.from_numpy(np.array(new)).to(int)
        return t.unsqueeze(0), c_text
    



class PredMaker():

    def __init__(self, model1, model2, rubert, model3, vectorizer, texter, clean_func, tokenizer, itc) -> None:
        self.log_reg_vec = model1
        self.lstm = model2
        self.bert = rubert
        self.log_reg_bert = model3

        self.vec = vectorizer
        self.clean = clean_func
        self.texter = texter
        self.ItV = itc
        self.tokenizer = tokenizer


        self.lstm.eval()
        self.bert.eval()

    def __call__(self, text):

        time1 = time()
        res1 = self.log_reg_vec.predict(self.vec.transform([self.clean(text)]))
        time1 = time() - time1

        time2 = time()
        t= self.texter(text)

        
        with torch.no_grad():
            res2 = torch.argmax(self.lstm(t[0])[0]).item()
        time2 = time() - time2

        tt = self.lstm(t[0])[1].detach().cpu().numpy()[0]


        time3 = time()
        text3 = self.tokenizer(text, truncation=True, return_tensors='pt',max_length=512)

        
        input = text3['input_ids']
        mask = text3['attention_mask']

        with torch.no_grad():
            
            embeddings = torch.nn.functional.normalize(self.bert(input.to(self.bert.device),mask.to(self.bert.device)).last_hidden_state[:, 0, :]).cpu().numpy()
        
        res3 = self.log_reg_bert.predict(embeddings)
        time3 = time() - time3
        
        return int(res1), int(res2), int(res3), t, tt, time1, time2, time3