"""Bag-of-words preprocessing for tiny sentiment classifier."""

from __future__ import annotations

import json
import re
from pathlib import Path

import numpy as np

VOCAB_PATH = Path(__file__).resolve().parent / "vocab.json"

TOKEN_PATTERN = re.compile(r"[a-z]+")


def tokenize(text: str) -> list[str]:
    return TOKEN_PATTERN.findall(text.lower())


def build_vocab(texts: list[str]) -> list[str]:
    words: set[str] = set()
    for text in texts:
        words.update(tokenize(text))
    return sorted(words)


def text_to_bow(text: str, vocab: list[str]) -> np.ndarray:
    word_to_idx = {word: idx for idx, word in enumerate(vocab)}
    bow = np.zeros(len(vocab), dtype=np.float32)
    for token in tokenize(text):
        idx = word_to_idx.get(token)
        if idx is not None:
            bow[idx] = 1.0
    return bow


def save_vocab(vocab: list[str], labels: list[str]) -> None:
    VOCAB_PATH.write_text(json.dumps({"vocab": vocab, "labels": labels}, indent=2))


def load_vocab() -> tuple[list[str], list[str]]:
    data = json.loads(VOCAB_PATH.read_text())
    return data["vocab"], data["labels"]