#!/usr/bin/env python3
"""Build bag-of-words and TF-IDF text baselines with the standard library.

The corpus is small and hand-written. Its labels are intentionally easy to
separate, so the reported score checks the implementation rather than showing
real-world NLP quality. Vocabulary and IDF statistics are fitted on training
documents only.
"""

from __future__ import annotations

import argparse
import math
import re
from collections import Counter
from collections.abc import Iterable, Mapping, Sequence


Document = tuple[str, str]
SparseVector = dict[int, float]

TRAIN_DOCUMENTS: tuple[Document, ...] = (
    ("space", "rocket launch reaches orbit around earth"),
    ("space", "astronaut trains for space station mission"),
    ("space", "satellite circles planet and sends images"),
    ("space", "telescope studies distant star and galaxy"),
    ("space", "spacecraft carries science instruments to moon"),
    ("cooking", "recipe mixes garlic tomato and herbs"),
    ("cooking", "chef bakes bread in a hot oven"),
    ("cooking", "kitchen knife chops onion and carrot"),
    ("cooking", "soup simmers with beans and spices"),
    ("cooking", "cook grills vegetables for dinner"),
    ("music", "guitar player performs a loud concert"),
    ("music", "piano melody fills the rehearsal room"),
    ("music", "drummer keeps rhythm for the band"),
    ("music", "singer records a new song in studio"),
    ("music", "violin musician practices classical music"),
)

VALIDATION_DOCUMENTS: tuple[Document, ...] = (
    ("space", "moon mission launches a rocket into orbit"),
    ("space", "telescope images a distant planet"),
    ("cooking", "chef chops tomato for soup"),
    ("cooking", "bread bakes in the kitchen oven"),
    ("music", "guitar and piano perform a song"),
    ("music", "the drummer plays rhythm at a concert"),
)

TEST_DOCUMENTS: tuple[Document, ...] = (
    ("space", "astronaut observes earth from the station"),
    ("space", "satellite mission studies the moon"),
    ("cooking", "a recipe uses vegetables and spices"),
    ("cooking", "cook onion and beans for dinner"),
    ("music", "the band records guitar music"),
    ("music", "singer practices a piano melody"),
)


def tokenize(text: str) -> list[str]:
    """Lowercase text and keep runs of Unicode word characters."""
    return re.findall(r"[^\W_]+(?:'[^\W_]+)?", text.lower(), flags=re.UNICODE)


def fit_vocabulary(texts: Iterable[str]) -> dict[str, int]:
    """Assign stable column IDs to tokens observed in training text."""
    terms = {token for text in texts for token in tokenize(text)}
    return {token: index for index, token in enumerate(sorted(terms))}


def vectorize_bow(
    text: str,
    vocabulary: Mapping[str, int],
    *,
    binary: bool = False,
) -> SparseVector:
    """Return a sparse count or presence vector; ignore out-of-vocabulary words."""
    counts = Counter(token for token in tokenize(text) if token in vocabulary)
    if binary:
        return {vocabulary[token]: 1.0 for token in counts}
    return {vocabulary[token]: float(count) for token, count in counts.items()}


def fit_idf(texts: Sequence[str], vocabulary: Mapping[str, int]) -> dict[int, float]:
    """Fit smoothed IDF: log((1 + N) / (1 + df)) + 1."""
    document_count = len(texts)
    if document_count == 0:
        raise ValueError("at least one training document is required")

    document_frequency: Counter[str] = Counter()
    for text in texts:
        document_frequency.update(set(tokenize(text)) & vocabulary.keys())

    return {
        index: math.log((1.0 + document_count) / (1.0 + document_frequency[token]))
        + 1.0
        for token, index in vocabulary.items()
    }


def vectorize_tfidf(
    text: str,
    vocabulary: Mapping[str, int],
    idf: Mapping[int, float],
) -> SparseVector:
    """Use raw token count as TF and multiply each non-zero value by IDF."""
    counts = vectorize_bow(text, vocabulary)
    return {index: count * idf[index] for index, count in counts.items()}


def cosine_similarity(left: Mapping[int, float], right: Mapping[int, float]) -> float:
    """Compute cosine similarity without constructing dense vocabulary-sized lists."""
    left_norm = math.sqrt(sum(value * value for value in left.values()))
    right_norm = math.sqrt(sum(value * value for value in right.values()))
    if left_norm == 0.0 or right_norm == 0.0:
        return 0.0
    if len(left) > len(right):
        left, right = right, left
    dot_product = sum(value * right.get(index, 0.0) for index, value in left.items())
    return dot_product / (left_norm * right_norm)


def mean_vector(vectors: Sequence[Mapping[int, float]]) -> SparseVector:
    """Return the element-wise arithmetic mean of sparse vectors."""
    if not vectors:
        raise ValueError("cannot average an empty vector collection")
    totals: Counter[int] = Counter()
    for vector in vectors:
        totals.update(vector)
    return {index: total / len(vectors) for index, total in totals.items()}


def fit_centroids(
    documents: Sequence[Document],
    vocabulary: Mapping[str, int],
    idf: Mapping[int, float],
) -> dict[str, SparseVector]:
    """Average training TF-IDF vectors separately for every class label."""
    grouped: dict[str, list[SparseVector]] = {}
    for label, text in documents:
        grouped.setdefault(label, []).append(vectorize_tfidf(text, vocabulary, idf))
    return {label: mean_vector(vectors) for label, vectors in grouped.items()}


def predict(
    text: str,
    vocabulary: Mapping[str, int],
    idf: Mapping[int, float],
    centroids: Mapping[str, Mapping[int, float]],
) -> tuple[str, dict[str, float]]:
    """Choose the label whose centroid has the highest cosine similarity."""
    query = vectorize_tfidf(text, vocabulary, idf)
    if not query:
        raise ValueError("text contains no words from the training vocabulary")
    scores = {
        label: cosine_similarity(query, centroid)
        for label, centroid in centroids.items()
    }
    label = max(sorted(scores), key=scores.get)
    return label, scores


def evaluate(
    documents: Sequence[Document],
    vocabulary: Mapping[str, int],
    idf: Mapping[int, float],
    centroids: Mapping[str, Mapping[int, float]],
) -> tuple[float, list[tuple[str, str, str]]]:
    """Return accuracy and rows of (expected, predicted, text)."""
    rows = []
    for expected, text in documents:
        predicted, _ = predict(text, vocabulary, idf, centroids)
        rows.append((expected, predicted, text))
    accuracy = sum(expected == predicted for expected, predicted, _ in rows) / len(rows)
    return accuracy, rows


def retrieve(
    query: str,
    documents: Sequence[Document],
    vocabulary: Mapping[str, int],
    idf: Mapping[int, float],
    *,
    limit: int = 3,
) -> list[tuple[float, str, str]]:
    """Rank documents by TF-IDF cosine similarity to a query."""
    if limit < 1:
        raise ValueError("limit must be positive")
    query_vector = vectorize_tfidf(query, vocabulary, idf)
    scored = [
        (
            cosine_similarity(query_vector, vectorize_tfidf(text, vocabulary, idf)),
            label,
            text,
        )
        for label, text in documents
    ]
    return sorted(scored, key=lambda item: (-item[0], item[2]))[:limit]


def run(*, binary_demo: bool = False) -> dict[str, float | int]:
    """Fit on training data, then report validation, test, and retrieval results."""
    training_texts = [text for _, text in TRAIN_DOCUMENTS]
    vocabulary = fit_vocabulary(training_texts)
    idf = fit_idf(training_texts, vocabulary)
    centroids = fit_centroids(TRAIN_DOCUMENTS, vocabulary, idf)

    validation_accuracy, _ = evaluate(
        VALIDATION_DOCUMENTS, vocabulary, idf, centroids
    )
    test_accuracy, test_rows = evaluate(TEST_DOCUMENTS, vocabulary, idf, centroids)
    majority_label_count = max(
        Counter(label for label, _ in TEST_DOCUMENTS).values()
    )
    majority_accuracy = majority_label_count / len(TEST_DOCUMENTS)

    sample = "rocket rocket moon unknownword"
    sample_vector = vectorize_bow(sample, vocabulary, binary=binary_demo)
    mode = "binary" if binary_demo else "count"
    print(f"training documents: {len(TRAIN_DOCUMENTS)}")
    print(f"training vocabulary size: {len(vocabulary)}")
    print(f"sample {mode} BoW non-zero values: {sample_vector}")
    print(f"majority-class test accuracy: {majority_accuracy:.3f}")
    print(f"validation accuracy: {validation_accuracy:.3f}")
    print(f"held-out test accuracy: {test_accuracy:.3f}")
    print("test predictions:")
    for expected, predicted, text in test_rows:
        print(f"  expected={expected:7s} predicted={predicted:7s} text={text}")

    print("retrieval for 'rocket moon mission':")
    for score, label, text in retrieve(
        "rocket moon mission", TRAIN_DOCUMENTS, vocabulary, idf
    ):
        print(f"  {score:.3f} [{label}] {text}")

    assert validation_accuracy >= 0.80
    assert test_accuracy >= 0.80
    assert test_accuracy > majority_accuracy
    return {
        "vocabulary_size": len(vocabulary),
        "majority_accuracy": majority_accuracy,
        "validation_accuracy": validation_accuracy,
        "test_accuracy": test_accuracy,
    }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--binary-demo",
        action="store_true",
        help="show word presence rather than word counts in the printed BoW example",
    )
    return parser.parse_args()


if __name__ == "__main__":
    arguments = parse_args()
    run(binary_demo=arguments.binary_demo)