---

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"<(?:user|assistant|eos|unk)>"
        r"|[а-яёa-z0-9]+(?:[-'][а-яёa-z0-9]+)*"
        r"|[.!?,;:…—–()«»\"“”]"
    )
    return re.findall(pattern, text.lower())


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


def format_tokens(tokens):
    """Склеивает токены в текст с аккуратными пробелами."""
    no_space_before = {
        ".", "!", "?", ",", ";", ":", "…", ")", "»"
    }
    no_space_after = {"(", "«"}

    result = ""

    for token in tokens:
        if token in {"<user>", "<assistant>", "<eos>", "<unk>"}:
            if result:
                result += " "
            result += token
        elif not result:
            result = token
        elif token in no_space_before or result.endswith(tuple(no_space_after)):
            result += token
        else:
            result += " " + token

    return result


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

    def __init__(self, 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),
        }

        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."""
        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 token in tokens:
        counts[token] = counts.get(token, 0) + 1

    # Не удаляем служебные маркеры даже при ограничении словаря.
    special_tokens = {"<user>", "<assistant>", "<eos>", "<unk>"}

    if len(counts) > MAX_VOCAB_SIZE:
        frequent = {
            token
            for token, _ in sorted(
                counts.items(),
                key=lambda item: item[1],
                reverse=True
            )[:MAX_VOCAB_SIZE]
        }
        frequent.update(special_tokens)

        tokens = [
            token if token in frequent else "<unk>"
            for token in tokens
        ]

    vocabulary = sorted(set(tokens))

    for token in ("<user>", "<assistant>", "<eos>", "<unk>"):
        if token not in vocabulary:
            vocabulary.append(token)

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

    token_ids = [word_to_id[token] for token 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):
    """Генерирует ответ на реплику пользователя."""
    if model is None:
        print("ИИ: Сначала введи /читать.")
        return

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

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

    generated = []

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

        temperature = max(float(temperature), 1e-3)
        adjusted = np.log(probabilities + 1e-12) / temperature
        probabilities = softmax(adjusted)

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

        if next_token == "<eos>":
            break

        # Не даём модели начинать новую реплику внутри ответа.
        if next_token in {"<user>", "<assistant>"}:
            continue

        context_ids.append(next_id)
        generated.append(next_token)

    answer = format_tokens(generated)
    print("ИИ:", answer if answer else "...")


def main():
    print("Мини-RNN запущена. Команды: /читать, /выход")
    print("Обучающие диалоги должны быть в book.txt.")
    print("Пример: <user> привет <assistant> привет! <eos>")

    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()