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

# Чтение текста из файла
with open('text_russian.txt', 'r', encoding='utf-8') as file:
    text = file.read()
# with open('text_english.txt', 'r', encoding='utf-8') as file:
#     text = file.read()

# Параметры модели
seq_length = 10

# Создание экземпляра Tokenizer и обучение на тексте
tokenizer = Tokenizer(char_level=True)
tokenizer.fit_on_texts([text])

# Преобразование текста в последовательности чисел
sequences = tokenizer.texts_to_sequences([text])[0]

# Создание входных и выходных последовательностей
X_data = []
y_data = []
for i in range(seq_length, len(sequences)):
    sequence = sequences[i - seq_length:i]
    target = sequences[i]
    X_data.append(sequence)
    y_data.append(target)

# Преобразование входных и выходных данных в формат массивов numpy
X = pad_sequences(X_data, maxlen=seq_length)
y = np.array(y_data)

# Создание модели RNN
vocab_size = len(tokenizer.word_index) + 1
model = Sequential()
model.add(LSTM(256, input_shape=(seq_length, 1), return_sequences=True))
model.add(LSTM(128, input_shape=(seq_length, 1)))
model.add(Dense(vocab_size, activation='softmax'))

# Компиляция модели
model.compile(loss='sparse_categorical_crossentropy', optimizer='adam', metrics=['accuracy'])

# Обучение модели
model.fit(X, y, epochs=100, verbose=1)


# Функция для генерации текста
def generate_text(seed_text, gen_length):
    generated_text = seed_text

    for _ in range(gen_length):
        sequence = tokenizer.texts_to_sequences([seed_text])[0]
        sequence = pad_sequences([sequence], maxlen=seq_length)
        prediction = model.predict(sequence)[0]
        predicted_index = np.argmax(prediction)
        predicted_char = tokenizer.index_word[predicted_index]
        generated_text += predicted_char
        seed_text += predicted_char
        seed_text = seed_text[1:]

    return generated_text

# Генерация текста
seed_text = "Уходили в тайгу"
# seed_text = "I had last night"
generated_text = generate_text(seed_text, 250)
print(generated_text)