import word_emb
from m_conf import *
import numpy as np
from gensim.models import Word2Vec
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Dropout, Flatten, Embedding
from keras_self_attention import SeqSelfAttention, SeqWeightedAttention
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.preprocessing.text import Tokenizer
from tensorflow.keras.losses import MeanSquaredError
from tensorflow.keras.preprocessing.sequence import pad_sequences

w2v = Word2Vec.load("w2v.model")

# загрузка датасета
with open('train.txt', 'r') as file:
    text = file.readlines()

# создание Tokenizerа
tokenizer = Tokenizer()
# обучение Tokenizer на основе текста из train.txt
tokenizer.fit_on_texts(text)

# преобразование текстовых данных в последовательности целых чисел с помощью объекта tokenizer
tt = tokenizer.texts_to_sequences(text)

t_sw = [[line[i:i+input_length] for i in range(len(line))] for line in tt]

combined_list = []

for line in t_sw:
    combined_list.extend(line)

y_t = [[w2v.wv[str(token)] for token in line] for line in tt]

y = []
for line in y_t:
    y.extend(line)

# задать длинну входа до переменной input_length, заполняя пустоту нулями
X = pad_sequences(combined_list, maxlen=input_length, padding='pre')

# получаем количество токенов в тексте
vocab_size = len(tokenizer.word_index)

# создание модели машинного обучения и задание её параметров
model = Sequential()
emb = Embedding(input_dim=vocab_size+1, output_dim=emb_dim, input_length=input_length)
model.add(emb)
model.add(SeqWeightedAttention())
model.add(Flatten())
model.add(Dense(512, activation="tanh"))
model.add(Dropout(0.5))
model.add(Dense(256, activation="tanh"))
model.add(Dropout(0.5))
model.add(Dense(128, activation="tanh"))
model.add(Dense(emb_o_dim, activation="tanh"))

# компилирование модели с функцией потерь mse и отображением accuracy
model.compile(optimizer=Adam(learning_rate=0.001), loss="mse", metrics=["accuracy"])

# обучение модели
set_limit = 2000
model.fit(np.array(X[:set_limit]), np.array(y[:set_limit]), epochs=10, batch_size=4)

def find_closest_token(o, temperature=0.0, top_p=1):
    token_distances = []
    for token in w2v.wv.index_to_key:
        vector = w2v.wv[token]
        distance = np.sum((o - vector)**2)
        token_distances.append((token, distance))

    token_distances = sorted(token_distances, key=lambda x: x[1])
    closest_token = token_distances[0][0]

    return closest_token

def gen(text):
    # преобразовать текст в понимаемую нейросетью информацию
    inp = pad_sequences(tokenizer.texts_to_sequences([text]), maxlen=input_length, padding='pre')
    # сделать предположение и его возвратить
    return str(tokenizer.index_word[int(find_closest_token(model.predict(inp)[0]))])