"""Week 6 demo 2: toy vector search using bag-of-words cosine similarity.

This mimics the retrieval idea without requiring an embedding model.
Real StatsChat uses sentence-transformer embeddings and FAISS.
"""

from collections import Counter
from math import sqrt

chunks = [
    "Annual inflation was 4.4 percent in March 2026 according to the CPI release.",
    "Mobile phone ownership was 53.7 percent nationally in the ICT factsheet.",
    "The housing survey describes household conditions and access to services.",
    "Food and transport prices contributed to changes in consumer prices.",
]
query = "Annual inflation in March 2026?"


STOPWORDS = {"what", "was", "in", "the", "and", "to", "of", "a", "an"}

def vectorise(text: str) -> Counter:
    words = text.lower().replace("?", "").replace(".", "").split()
    return Counter(w for w in words if w not in STOPWORDS)


def cosine(a: Counter, b: Counter) -> float:
    common = set(a) & set(b)
    dot = sum(a[w] * b[w] for w in common)
    norm_a = sqrt(sum(v * v for v in a.values()))
    norm_b = sqrt(sum(v * v for v in b.values()))
    return dot / (norm_a * norm_b) if norm_a and norm_b else 0.0

query_vec = vectorise(query)
ranked = sorted(
    ((cosine(query_vec, vectorise(chunk)), chunk) for chunk in chunks),
    reverse=True,
)

for score, chunk in ranked:
    print(f"{score:.2f}  {chunk}")