"""RAG for a tiny model, from scratch, readable.

A 27M story-writer does not KNOW facts — but it is a great COPIER.
So we give it memory the honest way:
  1. RETRIEVE: find the corpus line that shares the most words with the
     question (bag-of-words overlap — no embeddings needed at this scale).
  2. AUGMENT:  put that line into the prompt as context.
  3. GENERATE: let the model continue a cloze ("Mila the fox lives in ...") —
     with the fact in context it copies the truth; without it, it invents.

That difference (truth with context, guess without) is exactly what
contract A15 measures.
"""

import re
from pathlib import Path

import torch

STOP = {"the", "a", "an", "in", "on", "under", "by", "to", "of", "and",
        "does", "did", "do", "what", "where", "when", "who", "how",
        "is", "are", "was", "were", "her", "his", "its"}

_word_re = re.compile(r"\S+\s*")
_cache = {}


def fast_encode(tok, text):
    """Word-cache encoder — the SAME one every training stage used.
    Encoding any prompt differently feeds the model unfamiliar token
    sequences and its predictions turn to mush (M6 lesson: we tried)."""
    out = []
    for w in _word_re.findall(text):
        ids = _cache.get(w)
        if ids is None:
            ids = tok.encode(w)
            _cache[w] = ids
        out.extend(ids)
    return out


def words(text):
    return [w for w in re.findall(r"[a-z]+", text.lower()) if w not in STOP]


def load_corpus(path):
    return [ln.strip() for ln in Path(path).read_text().splitlines() if ln.strip()]


def retrieve(question, corpus, k=2):
    """Top-k corpus lines by shared-word count with the question."""
    q = set(words(question))
    scored = sorted(corpus, key=lambda ln: -len(q & set(words(ln))))
    return scored[:k]


@torch.no_grad()
def complete(model, tok, prompt, n=14):
    """Greedy continuation — deterministic, so the gate is reproducible."""
    ids = torch.tensor([fast_encode(tok, prompt)])
    n_prompt = ids.shape[1]
    for _ in range(n):
        logits, _ = model(ids[:, -model.cfg.max_seq_len:])
        nxt = logits[:, -1, :].argmax(-1, keepdim=True)
        ids = torch.cat([ids, nxt], dim=1)
    return tok.decode(ids[0, n_prompt:].tolist())


def answer(model, tok, question, cloze, corpus, use_rag=True):
    """Answer a fact question by continuing its cloze, with/without context."""
    if use_rag:
        ctx = " ".join(retrieve(question, corpus))
        prompt = f"{ctx}\n\n{cloze}"
    else:
        prompt = cloze
    return complete(model, tok, prompt)


if __name__ == "__main__":
    import sys
    sys.path.insert(0, str(Path(__file__).parent))
    from bpe import BPETokenizer
    from model import TinyLLM, ModelConfig

    ROOT = Path(__file__).resolve().parent.parent
    tok = BPETokenizer.load(str(ROOT / "data" / "tokenizer.json"))
    ck = torch.load(ROOT / "checkpoints" / "dpo-m4c.pt", map_location="cpu")
    m = TinyLLM(ModelConfig(**ck["cfg"])); m.load_state_dict(ck["model"]); m.eval()
    corpus = load_corpus(ROOT / "data" / "rag-corpus.txt")

    q, cloze = "Where does Mila the fox live?", "Mila the fox lives in"
    print("retrieved:", retrieve(q, corpus)[0])
    print("WITH rag   :", answer(m, tok, q, cloze, corpus, True).strip())
    print("WITHOUT rag:", answer(m, tok, q, cloze, corpus, False).strip())