from datasets import load_dataset
from collections import Counter
import torch
import torch.nn as nn
import torch.optim as optim
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DataLoader
import random
import math
import pandas as pd
from tabulate import tabulate

import re

import time
from tqdm import tqdm  # для отображения прогресс-бара

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Используется устройство: {device}")

# Загружаем IMDB
dataset = load_dataset('imdb')

# Простая токенизация без NLTK
def tokenize(text):
    # Оставляем только буквы и апострофы (для "don't" и т.п.)
    return re.findall(r"\b\w+\b", text.lower())

# Собираем все слова из обучающей выборки
all_words = []
for example in dataset['train']:
    all_words.extend(tokenize(example['text']))

# обучающий пример
print(dataset['train'][3])

# Зададим, сколько самых частых слов из корпуса мы хотим включить в словарь.
# Это гиперпараметр: можно уменьшить (для быстрого обучения) или увеличить (для лучшего покрытия).
vocab_size = 10000

# Counter подсчитывает, сколько раз каждое слово встречается в обучающем корпусе.
counter = Counter(all_words)

# Служебные токены
#<pad> - Заполнение коротких последовательностей до одинаковой длины (для батчей)
# <unk> Замена редких/неизвестных слов (out-of-vocabulary)
# <bos> Beginning of sentence — маркер начала предложения
# <eos> End of sentence — маркер конца предложения
special_tokens = ['<pad>', '<unk>', '<bos>', '<eos>']

# Берём 10 000 самых частых слов из корпуса (без учёта служебных).
# Редкие слова автоматически будут заменяться на <unk>.
most_common = [word for word, _ in counter.most_common(vocab_size)]

# Формируем итоговый словарь: сначала идут служебные токены, затем частые слова.
vocab = special_tokens + most_common
word2idx = {word: idx for idx, word in enumerate(vocab)}
idx2word = {idx: word for word, idx in word2idx.items()}

print(f"Размер словаря: {len(vocab)}")

"""Данная функция преобразует текстовую строку в последовательность числовых индексов,
пригодную для подачи в нейронную сеть (например, RNN).
"""

def text_to_indices(text):
    tokens = tokenize(text)
    # добавление начала предложения
    indices = [word2idx['<bos>']]
    for token in tokens:
        indices.append(word2idx.get(token, word2idx['<unk>']))

    # добавление конца предложения
    indices.append(word2idx['<eos>'])
    return torch.tensor(indices, dtype=torch.long)

"""Данная функция подготавливает подмножество данных (раздел датасета)
для обучения или тестирования языковой модели."""
def prepare_split(split_name, max_samples=5000):
    data = []
    for i, example in enumerate(dataset[split_name]):
        if i >= max_samples:
            break
        data.append(text_to_indices(example['text']))
    return data

train_data = prepare_split('train', max_samples=8000)
test_data = prepare_split('test', max_samples=2000)

train_data[0]

class RNNLanguageModel(nn.Module):
    def __init__(self, vocab_size, embed_dim=200, hidden_dim=200,
                 num_layers=1, rnn_type='LSTM', dropout=0.0):
        super().__init__()
        # Эмбеддинг-слой: преобразует индексы слов в плотные векторы
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=word2idx['<pad>'])
        # Dropout после эмбеддингов для регуляризации
        self.dropout_emb = nn.Dropout(dropout)

        # Выбор типа рекуррентного слоя
        rnn_types = {
            'RNN': nn.RNN,
            'GRU': nn.GRU,
            'LSTM': nn.LSTM
        }
        # Инициализация RNN-блока
        self.rnn = rnn_types[rnn_type](
            input_size=embed_dim,      # размер входного эмбеддинга
            hidden_size=hidden_dim,    # размер скрытого состояния
            num_layers=num_layers,     # количество слоёв
            dropout=dropout if num_layers > 1 else 0,  # dropout между слоями (только если их >1)
            batch_first=False          # ожидаем форму (seq_len, batch, ...)
        )
        # Dropout на выходе RNN
        self.dropout_out = nn.Dropout(dropout)
        # Линейный слой: преобразует скрытое состояние в логиты над словарём
        self.fc = nn.Linear(hidden_dim, vocab_size)
        # Сохраняем тип RNN для возможного использования в forward
        self.rnn_type = rnn_type

    def forward(self, x, hidden=None):
        # x: тензор формы (seq_len, batch_size) с индексами слов
        emb = self.embedding(x)               # (seq_len, batch, embed_dim)
        emb = self.dropout_emb(emb)           # применяем dropout к эмбеддингам

        # Пропускаем через RNN
        output, hidden = self.rnn(emb, hidden)  # output: (seq_len, batch, hidden_dim)

        output = self.dropout_out(output)      # dropout на выходе RNN
        logits = self.fc(output)              # (seq_len, batch, vocab_size)

        return logits, hidden

# Преобразуем список последовательностей разной длины в единый тензор фиксированной формы, который
#можно подать в нейронную сеть.
def collate_batch(batch):
    batch = [seq[:200] for seq in batch if len(seq) > 1]  # ограничиваем длину и фильтруем короткие
    padded = pad_sequence(batch, batch_first=False, padding_value=word2idx['<pad>'])
    return padded.to(device)

# Оценка модели на тестовых данных: вычисляет среднюю кросс-энтропию (loss)
def evaluate(model, data_loader, criterion):
    model.eval()
    total_loss = 0.0
    total_tokens = 0
    with torch.no_grad():  # отключаем градиенты для ускорения
        for batch in data_loader:
            if batch.size(0) <= 1:
                continue
            x = batch[:-1, :]          # вход: все слова, кроме последнего
            y = batch[1:, :].reshape(-1)  # цель: все слова, кроме первого
            logits, _ = model(x)
            loss = criterion(logits.reshape(-1, len(vocab)), y)
            total_loss += loss.item() * y.numel()
            total_tokens += y.numel()
    return total_loss / total_tokens  # средний loss на токен

# Обучение модели
def train_model(model, train_loader, epochs=3, lr=0.002, clip=0.5):
    # Переносим модель на GPU (если доступен) или CPU
    model.to(device)

    # Функция потерь: кросс-энтропия, игнорирующая паддинг-токены при подсчёте
    criterion = nn.CrossEntropyLoss(ignore_index=word2idx['<pad>'])

    # Используем только Adam — адаптивный оптимизатор, устойчивый к выбору lr
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)

    # Планировщик: уменьшает learning rate на 5% после каждой эпохи для стабильной сходимости
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.95)

    for epoch in range(epochs):
        model.train()  # включаем режим обучения (включает dropout и т.п.)
        total_loss = 0
        start_time = time.time()

        # Прогресс-бар для отслеживания обучения внутри эпохи
        progress_bar = tqdm(train_loader, desc=f"Эпоха {epoch+1}/{epochs}", leave=False)
        for batch in progress_bar:
            # Пропускаем слишком короткие последовательности (нельзя предсказать следующее слово)
            if batch.size(0) <= 1:
                continue

            # x — вход: все слова, кроме последнего (<eos>)
            # y — цель: все слова, кроме первого (<bos>)
            x = batch[:-1, :]
            y = batch[1:, :].reshape(-1)

            optimizer.zero_grad()          # обнуляем градиенты с прошлой итерации
            logits, _ = model(x)           # прямой проход: получаем логиты
            loss = criterion(logits.reshape(-1, len(vocab)), y)  # вычисляем потерю
            loss.backward()                # обратное распространение ошибки
            torch.nn.utils.clip_grad_norm_(model.parameters(), clip)  # обрезка градиентов против взрыва
            optimizer.step()               # обновление весов

            total_loss += loss.item()
            progress_bar.set_postfix({'loss': f"{loss.item():.3f}"})  # отображаем текущий loss

        # Средний loss за эпоху и затраченное время
        avg_loss = total_loss / len(train_loader)
        elapsed = time.time() - start_time
        print(f"Эпоха {epoch+1} завершена. Средний loss: {avg_loss:.3f}. Время: {elapsed:.1f} сек.")

        # Уменьшаем learning rate
        scheduler.step()

    return model

experiments = []

# Подготовка данных
train_loader = DataLoader(train_data, batch_size=32, shuffle=True, collate_fn=collate_batch)
test_loader = DataLoader(test_data, batch_size=32, shuffle=False, collate_fn=collate_batch)

print("\n🚀 Эксперимент 3: Размер скрытого состояния")
for hidden_dim in [50, 100, 200]:
    print(f"  Обучение LSTM, hidden_dim={hidden_dim}...")
    model = RNNLanguageModel(
        vocab_size=len(vocab),
        embed_dim=hidden_dim,  # обычно embed_dim = hidden_dim для простоты
        hidden_dim=hidden_dim,
        num_layers=1,
        rnn_type='LSTM',
        dropout=0.2
    )
    model = train_model(model, train_loader, test_loader, epochs=3, lr=0.002)

    criterion = nn.CrossEntropyLoss(ignore_index=word2idx['<pad>'])
    test_loss = evaluate(model, test_loader, criterion)
    ppl = math.exp(test_loss)
    experiments.append({
        'Модель': f'LSTM (hidden={hidden_dim})',
        'Скрытый размер': hidden_dim,
        'Слои': 1,
        'Dropout': 0.2,
        'Оптимизатор': 'Adam',
        'Перплексия': round(ppl, 1)
    })

df = pd.DataFrame(experiments)
print("\n=== Результаты экспериментов ===")
print(tabulate(df, headers='keys', tablefmt='grid', showindex=False))

def generate_text(model, seed_text, max_len=40, temperature=1.0, method='greedy'):
    """
    Генерирует текст с помощью обученной RNN-модели.

    Аргументы:
        model: обученная RNNLanguageModel
        seed_text (str): начальный текст (например, "I love")
        max_len (int): максимальная длина генерируемого текста (в словах)
        temperature (float): температура для sampling'а (1.0 = стандартная)
        method (str): 'greedy' или 'sample'
    """
    model.eval()
    tokens = tokenize(seed_text)
    indices = [word2idx['<bos>']] + [word2idx.get(t, word2idx['<unk>']) for t in tokens]
    generated = tokens[:]

    with torch.no_grad():
        for _ in range(max_len):
            # Подготавливаем входной тензор
            x = torch.tensor(indices, dtype=torch.long).unsqueeze(1).to(device)  # (seq_len, 1)

            # Прямой проход
            logits, _ = model(x)

            # Берём логиты последнего слова
            next_logits = logits[-1, 0, :] / temperature
            probs = torch.softmax(next_logits, dim=-1)

            if method == 'greedy':
                next_idx = torch.argmax(probs).item()
            else:  # sampling
                next_idx = torch.multinomial(probs, 1).item()

            # Остановка при <eos>
            if next_idx == word2idx['<eos>']:
                break

            # Добавляем слово
            word = idx2word.get(next_idx, '<unk>')
            generated.append(word)
            indices.append(next_idx)

    return ' '.join(generated)

# После обучения модели (например, lstm_model)
print("=== Примеры генерации ===")
seeds = ["i", "i think", "i think this", "i think this film"]

for seed in seeds:
    greedy = generate_text(model, seed, max_len=25, method='greedy')
    sampled = generate_text(model, seed, max_len=25, method='sample', temperature=0.8)
    print(f"Начало: '{seed}'")
    print(f"  Жадно:     {greedy}")
    print(f"  Sampling:  {sampled}\n")