import numpy as np
import matplotlib.pyplot as plt
from tensorflow.keras.datasets import imdb
from tensorflow.keras.preprocessing.sequence import pad_sequences
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Embedding, SimpleRNN, GRU, LSTM, Dense, Dropout
from tensorflow.keras.optimizers import Adam
from sklearn.metrics import classification_report, confusion_matrix

# --- 1. Загрузка и предобработка данных ---

# Гиперпараметры
VOCAB_SIZE = 10000  # Размер словаря (10,000 наиболее частых слов)
MAX_LEN = 256  # Максимальная длина последовательности
EMBEDDING_DIM = 128  # Размерность векторных представлений слов
BATCH_SIZE = 64
EPOCHS = 5

# Загрузка датасета IMDB
print("Загрузка данных...")
(x_train, y_train), (x_test, y_test) = imdb.load_data(num_words=VOCAB_SIZE)
print(f"Обучающая выборка: {len(x_train)} отзывов")
print(f"Тестовая выборка: {len(x_test)} отзывов")

# Приведение последовательностей к одной длине
print(f"Паддинг/обрезка последовательностей до длины {MAX_LEN}...")
x_train_padded = pad_sequences(x_train, maxlen=MAX_LEN, padding='post', truncating='post')
x_test_padded = pad_sequences(x_test, maxlen=MAX_LEN, padding='post', truncating='post')
print(f"Размер обучающей выборки после паддинга: {x_train_padded.shape}")
print(f"Размер тестовой выборки после паддинга: {x_test_padded.shape}")



# --- 2. Создание модели RNN ---

def build_model(rnn_type='LSTM', embedding_dim=128, rnn_units=64, learning_rate=0.001):
    """
    Функция для построения модели с выбором типа RNN слоя.
    """
    print(f"\nСоздание модели с {rnn_type} слоем...")
    model = Sequential()

    # Слой эмбеддингов
    model.add(Embedding(input_dim=VOCAB_SIZE, output_dim=embedding_dim, input_length=MAX_LEN))

    # Слой Dropout для регуляризации эмбеддингов
    model.add(Dropout(0.5))

    # RNN слой (SimpleRNN, GRU или LSTM)
    if rnn_type == 'SimpleRNN':
        model.add(SimpleRNN(rnn_units))
    elif rnn_type == 'GRU':
        model.add(GRU(rnn_units))
    else:  # LSTM по умолчанию
        model.add(LSTM(rnn_units))

    # Слой Dropout для регуляризации RNN слоя
    model.add(Dropout(0.5))

    # Выходной слой для бинарной классификации
    model.add(Dense(1, activation='sigmoid'))

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

    return model


# Создадим модель (например, с LSTM)
model = build_model(rnn_type='LSTM', embedding_dim=EMBEDDING_DIM)
model.summary()

# --- 3. Обучение модели ---

print("\nНачало обучения модели...")
history = model.fit(
    x_train_padded, y_train,
    epochs=EPOCHS,
    batch_size=BATCH_SIZE,
    validation_split=0.2  # Используем 20% обучающих данных для валидации
)

# --- 4. Оценка качества ---

# Оценка на тестовых данных
print("\nОценка модели на тестовых данных...")
loss, accuracy = model.evaluate(x_test_padded, y_test, verbose=0)
print(f"Точность (Accuracy) на тесте: {accuracy:.4f}")
print(f"Потери (Loss) на тесте: {loss:.4f}")

# Получение предсказаний
y_pred_prob = model.predict(x_test_padded)
y_pred = (y_pred_prob > 0.5).astype("int32")

# Расчет Precision, Recall, F1-score
print("\nОтчет по классификации:")
print(classification_report(y_test, y_pred, target_names=['Negative', 'Positive']))


# Построение графиков loss и accuracy
def plot_history(history):
    acc = history.history['accuracy']
    val_acc = history.history['val_accuracy']
    loss = history.history['loss']
    val_loss = history.history['val_loss']
    epochs_range = range(1, len(acc) + 1)

    plt.figure(figsize=(12, 5))

    plt.subplot(1, 2, 1)
    plt.plot(epochs_range, acc, 'bo-', label='Training Acc')
    plt.plot(epochs_range, val_acc, 'ro-', label='Validation Acc')
    plt.title('Точность (Accuracy) на обучении и валидации')
    plt.xlabel('Эпохи')
    plt.ylabel('Точность')
    plt.legend()

    plt.subplot(1, 2, 2)
    plt.plot(epochs_range, loss, 'bo-', label='Training Loss')
    plt.plot(epochs_range, val_loss, 'ro-', label='Validation Loss')
    plt.title('Потери (Loss) на обучении и валидации')
    plt.xlabel('Эпохи')
    plt.ylabel('Потери')
    plt.legend()

    plt.show()


plot_history(history)


# Анализ ошибок
def analyze_errors():
    # Загружаем словарь для декодирования
    word_index = imdb.get_word_index()
    reverse_word_index = {v: k for k, v in word_index.items()}

    def decode_review(text_indices):
        # Смещения стандартные для датасета Keras
        return ' '.join([reverse_word_index.get(i - 3, '?') for i in text_indices])

    print("\nАнализ неверно классифицированных отзывов:")
    misclassified_indices = np.where(y_pred.flatten() != y_test)[0]

    for i in range(5):  # Посмотрим на 5 случайных ошибок
        if i < len(misclassified_indices):
            idx = misclassified_indices[i]
            print(f"\n--- Отзыв #{idx} ---")
            print(f"Текст: {decode_review(x_test[idx])}")
            print(f"Настоящая метка: {'Positive' if y_test[idx] == 1 else 'Negative'}")
            print(f"Предсказанная метка: {'Positive' if y_pred[idx] == 1 else 'Negative'}")


analyze_errors()