---

import re
import numpy as np


BOOK_FILE = "book.txt"

EMBED_SIZE = 32
HIDDEN_SIZE = 64
SEQUENCE_LENGTH = 24
EPOCHS = 20
LEARNING_RATE = 0.003

# Ограничивает размер словаря на очень больших книгах.
MAX_VOCAB_SIZE = 5000

rng = np.random.default_rng(42)

vocabulary = []
word_to_id = {}
model = None


def tokenize(text):
    """Разбивает текст на слова, числа и знаки пунктуации."""
    pattern = (
        r"<unk>"
        r"|[а-яёa-z0-9]+(?:[-'][а-яёa-z0-9]+)*"
        r"|[.!?,;:…—–()«»\"“”]"
    )
    return re.findall(pattern, text.lower())


def format_tokens(tokens):
    """Склеивает токены в текст с более естественными пробелами."""
    if not tokens:
        return ""

    no_space_before = {".", "!", "?", ",", ";", ":", "…", ")", "»"}
    no_space_after = {"(", "«"}

    result = tokens[0]

    for token in tokens[1:]:
        if token in no_space_before:
            result += token
        elif result[-1:] in no_space_after:
            result += token
        else:
            result += " " + token

    return result


def softmax(values):
    values = values - np.max(values)
    exp_values = np.exp(values)
    return exp_values / np.sum(exp_values)


class WordRNN:
    """Однослойная RNN: embedding -> tanh RNN -> softmax."""

    def __init__(self, vocab_size):
        self.vocab_size = vocab_size

        self.params = {
            "E": rng.normal(0, 0.08, (vocab_size, EMBED_SIZE)),
            "Wxh": rng.normal(
                0, 1 / np.sqrt(EMBED_SIZE), (HIDDEN_SIZE, EMBED_SIZE)
            ),
            "Whh": rng.normal(
                0, 1 / np.sqrt(HIDDEN_SIZE), (HIDDEN_SIZE, HIDDEN_SIZE)
            ),
            "bh": np.zeros(HIDDEN_SIZE),
            "Why": rng.normal(
                0, 1 / np.sqrt(HIDDEN_SIZE), (vocab_size, HIDDEN_SIZE)
            ),
            "by": np.zeros(vocab_size),
        }

        # Состояние Adam для каждого обучаемого параметра.
        self.m = {
            name: np.zeros_like(value)
            for name, value in self.params.items()
        }
        self.v = {
            name: np.zeros_like(value)
            for name, value in self.params.items()
        }
        self.step = 0

    def loss_and_gradients(self, inputs, targets):
        """Считает средний Cross-Entropy Loss и градиенты через BPTT."""
        if not inputs:
            return 0.0, {
                name: np.zeros_like(value)
                for name, value in self.params.items()
            }

        params = self.params
        hidden_states = [np.zeros(HIDDEN_SIZE)]
        probabilities = []
        loss = 0.0

        # Прямой проход.
        for word_id, target_id in zip(inputs, targets):
            x = params["E"][word_id]
            previous_hidden = hidden_states[-1]

            hidden = np.tanh(
                params["Wxh"] @ x
                + params["Whh"] @ previous_hidden
                + params["bh"]
            )

            logits = params["Why"] @ hidden + params["by"]
            probs = softmax(logits)

            loss -= np.log(probs[target_id] + 1e-12)
            hidden_states.append(hidden)
            probabilities.append(probs)

        length = len(inputs)
        loss /= length

        grads = {
            name: np.zeros_like(value)
            for name, value in params.items()
        }
        next_hidden_gradient = np.zeros(HIDDEN_SIZE)

        # Обратный проход по времени.
        for t in reversed(range(length)):
            target_id = targets[t]
            hidden = hidden_states[t + 1]
            previous_hidden = hidden_states[t]

            dlogits = probabilities[t].copy()
            dlogits[target_id] -= 1.0
            dlogits /= length

            grads["Why"] += np.outer(dlogits, hidden)
            grads["by"] += dlogits

            dhidden = params["Why"].T @ dlogits + next_hidden_gradient
            dpreactivation = dhidden * (1.0 - hidden * hidden)

            grads["Wxh"] += np.outer(
                dpreactivation,
                params["E"][inputs[t]]
            )
            grads["Whh"] += np.outer(
                dpreactivation,
                previous_hidden
            )
            grads["bh"] += dpreactivation
            grads["E"][inputs[t]] += (
                params["Wxh"].T @ dpreactivation
            )

            next_hidden_gradient = (
                params["Whh"].T @ dpreactivation
            )

        return loss, grads

    def update_adam(self, grads):
        """Обновляет параметры вручную, без autograd."""
        self.step += 1
        beta1 = 0.9
        beta2 = 0.999
        epsilon = 1e-8

        for name in self.params:
            grad = np.clip(grads[name], -5.0, 5.0)

            self.m[name] = (
                beta1 * self.m[name]
                + (1 - beta1) * grad
            )
            self.v[name] = (
                beta2 * self.v[name]
                + (1 - beta2) * (grad * grad)
            )

            m_hat = self.m[name] / (1 - beta1 ** self.step)
            v_hat = self.v[name] / (1 - beta2 ** self.step)

            self.params[name] -= (
                LEARNING_RATE
                * m_hat
                / (np.sqrt(v_hat) + epsilon)
            )

    def train(self, token_ids, epochs):
        """Обучает модель предсказывать следующее слово."""
        if len(token_ids) < 2:
            print("ИИ: Недостаточно слов для обучения.")
            return

        starts = list(range(0, len(token_ids) - 1, SEQUENCE_LENGTH))

        for epoch in range(epochs):
            total_loss = 0.0
            total_tokens = 0

            for start in starts:
                end = min(start + SEQUENCE_LENGTH, len(token_ids) - 1)
                inputs = token_ids[start:end]
                targets = token_ids[start + 1:end + 1]

                if not inputs:
                    continue

                loss, grads = self.loss_and_gradients(inputs, targets)
                self.update_adam(grads)

                total_loss += loss * len(inputs)
                total_tokens += len(inputs)

            average_loss = total_loss / max(total_tokens, 1)
            print(
                f"Эпоха {epoch + 1}/{epochs} — "
                f"loss: {average_loss:.4f}"
            )

    def next_word_probabilities(self, context_ids):
        """Возвращает вероятности следующего слова."""
        hidden = np.zeros(HIDDEN_SIZE)

        for word_id in context_ids[-SEQUENCE_LENGTH:]:
            x = self.params["E"][word_id]
            hidden = np.tanh(
                self.params["Wxh"] @ x
                + self.params["Whh"] @ hidden
                + self.params["bh"]
            )

        logits = self.params["Why"] @ hidden + self.params["by"]
        return softmax(logits)


def read_book():
    """Загружает книгу, строит словарь и запускает обучение."""
    global vocabulary, word_to_id, model

    try:
        with open(BOOK_FILE, "r", encoding="utf-8") as file:
            text = file.read()
    except FileNotFoundError:
        print(
            f"ИИ: Не найден файл {BOOK_FILE}. "
            "Положи его рядом со скриптом."
        )
        return
    except OSError as error:
        print("ИИ: Не удалось прочитать книгу:", error)
        return

    tokens = tokenize(text)

    if len(tokens) < 3:
        print("ИИ: В книге слишком мало текста.")
        return

    counts = {}
    for word in tokens:
        counts[word] = counts.get(word, 0) + 1

    # Оставляем самые частые слова, редкие заменяем на <unk>.
    if len(counts) > MAX_VOCAB_SIZE:
        frequent_words = sorted(
            counts.items(),
            key=lambda item: item[1],
            reverse=True
        )[:MAX_VOCAB_SIZE - 1]

        frequent = {word for word, _ in frequent_words}
        tokens = [
            word if word in frequent else "<unk>"
            for word in tokens
        ]

    vocabulary = sorted(set(tokens))

    if "<unk>" not in vocabulary:
        vocabulary.append("<unk>")

    word_to_id = {
        word: index
        for index, word in enumerate(vocabulary)
    }

    token_ids = [
        word_to_id.get(word, word_to_id["<unk>"])
        for word in tokens
    ]

    model = WordRNN(len(vocabulary))

    print(
        f"ИИ: Прочитал токенов: {len(tokens)}; "
        f"слов в словаре: {len(vocabulary)}."
    )

    model.train(token_ids, EPOCHS)
    print("ИИ: Обучение завершено.")


def generate(prompt, max_words=40, temperature=0.8):
    """Продолжает фразу, выбирая слова из softmax-распределения."""
    if model is None:
        print("ИИ: Сначала введи /читать.")
        return

    prompt_tokens = tokenize(prompt)
    unk_id = word_to_id["<unk>"]

    context_ids = [
        word_to_id.get(word, unk_id)
        for word in prompt_tokens
    ]

    if not context_ids:
        context_ids = [int(rng.integers(len(vocabulary)))]

    generated = prompt_tokens[:]
    temperature = max(float(temperature), 1e-3)

    for _ in range(max_words):
        probabilities = model.next_word_probabilities(context_ids)

        # Температура управляет случайностью выбора.
        adjusted = np.log(probabilities + 1e-12) / temperature
        probabilities = softmax(adjusted)

        next_id = int(rng.choice(len(vocabulary), p=probabilities))
        next_word = vocabulary[next_id]

        context_ids.append(next_id)
        generated.append(next_word)

        if next_word in {".", "!", "?"}:
            break

    print("ИИ:", format_tokens(generated))


def main():
    print("Мини-RNN запущена. Команды: /читать, /выход")
    print("Сначала введи /читать, чтобы обучить модель на book.txt.")

    while True:
        try:
            text = input("Ты: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nИИ: Завершение работы.")
            break

        if text.lower() == "/выход":
            print("ИИ: Завершение работы.")
            break

        if text.lower() == "/читать":
            read_book()
            continue

        if not text:
            continue

        if model is None:
            print("ИИ: Сначала введи /читать.")
            continue

        generate(text)


if __name__ == "__main__":
    main()