import numpy as np
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Embedding, LSTM, Dense
from tensorflow.keras.preprocessing.text import Tokenizer
from tensorflow.keras.preprocessing.sequence import pad_sequences

# Загрузка русского текста
with open('text/rus.txt', encoding='utf-8') as file:
    text_rus = file.read()

# Загрузка английского текста
with open('text/eng.txt', encoding='utf-8') as file:
    text_eng = file.read()

# Токенизация русского текста
tokenizer_rus = Tokenizer()
tokenizer_rus.fit_on_texts([text_rus])
total_words_rus = len(tokenizer_rus.word_index) + 1

# Токенизация английского текста
tokenizer_eng = Tokenizer()
tokenizer_eng.fit_on_texts([text_eng])
total_words_eng = len(tokenizer_eng.word_index) + 1

# Создание последовательностей для обучения модели русского текста
input_sequences_rus = []
for line in text_rus.split('\n'):
    token_list = tokenizer_rus.texts_to_sequences([line])[0]
    for i in range(1, len(token_list)):
        n_gram_sequence = token_list[:i + 1]
        input_sequences_rus.append(n_gram_sequence)

max_sequence_length_rus = max([len(x) for x in input_sequences_rus])
input_sequences_rus = pad_sequences(input_sequences_rus, maxlen=max_sequence_length_rus, padding='pre')

predictors_rus, labels_rus = input_sequences_rus[:, :-1], input_sequences_rus[:, -1]

# Создание последовательностей для обучения модели английского текста
input_sequences_eng = []
for line in text_eng.split('\n'):
    token_list = tokenizer_eng.texts_to_sequences([line])[0]
    for i in range(1, len(token_list)):
        n_gram_sequence = token_list[:i + 1]
        input_sequences_eng.append(n_gram_sequence)

max_sequence_length_eng = max([len(x) for x in input_sequences_eng])
input_sequences_eng = pad_sequences(input_sequences_eng, maxlen=max_sequence_length_eng, padding='pre')

predictors_eng, labels_eng = input_sequences_eng[:, :-1], input_sequences_eng[:, -1]

# Создание и компиляция модели для русского текста
model_rus = Sequential([
    Embedding(total_words_rus, 100, input_length=max_sequence_length_rus - 1),
    LSTM(150),
    Dense(total_words_rus, activation='softmax')
])

model_rus.compile(loss='sparse_categorical_crossentropy', optimizer='adam', metrics=['accuracy'])

# Тренировка модели русского текста
model_rus.fit(predictors_rus, labels_rus, epochs=100, verbose=1)

# Создание и компиляция модели для английского текста
model_eng = Sequential([
    Embedding(total_words_eng, 100, input_length=max_sequence_length_eng - 1),
    LSTM(150),
    Dense(total_words_eng, activation='softmax')
])

model_eng.compile(loss='sparse_categorical_crossentropy', optimizer='adam', metrics=['accuracy'])

# Тренировка модели английского текста
model_eng.fit(predictors_eng, labels_eng, epochs=100, verbose=1)

# Генерация текста на основе русского текста
def generate_text(seed_text, next_words, model, max_sequence_length, tokenizer):
    for _ in range(next_words):
        token_list = tokenizer.texts_to_sequences([seed_text])[0]
        token_list = pad_sequences([token_list], maxlen=max_sequence_length - 1, padding='pre')
        predicted_probs = model.predict(token_list)[0]
        predicted_word_index = np.random.choice(len(predicted_probs), p=predicted_probs)
        output_word = tokenizer.index_word.get(predicted_word_index, '')
        seed_text += " " + output_word
    return seed_text

# Генерация текста на основе русского текста
generated_text_rus = generate_text("Начало", 50, model_rus, max_sequence_length_rus, tokenizer_rus)
print("Generated Russian Text:")
print(generated_text_rus)

# Генерация текста на основе английского текста
generated_text_eng = generate_text("Start", 50, model_eng, max_sequence_length_eng, tokenizer_eng)
print("\nGenerated English Text:")
print(generated_text_eng)