"""rag_toolkit.py - Week 2b lab: a tiny RAG pipeline. Standard library only.

Parts:
  1. Chunkers: fixed, fixed with overlap, sentence groups, contextual prefix
  2. Retrievers: BM25 from scratch, a toy dense retriever, hybrid (RRF)
  3. Reranker: a toy late-interaction (MaxSim) scorer
  4. Metrics: hit@k, recall@k, MRR, nDCG@k
  5. Generator: a mock LLM that answers only from sources, cites them, or abstains
  6. Agent loop: decide whether to retrieve, check the result, retry or stop

IMPORTANT: the "dense" retriever is a STAND-IN. It hashes words and letter
triples into a small vector and uses a hand-made synonym table. A real system
uses a trained neural embedding model. Nothing here needs an API key.

Run `python rag_toolkit.py` for a demo. Run `python test_rag_toolkit.py` for tests.
"""
import math
import re
import zlib
from dataclasses import dataclass, field

# ------------------------------------------------------------ text basics

STOPWORDS = set("""a an the of to in on at for and or is are was were be been do does did can could
may must should will would i you we they it my our your their me what which who when where how why
this that these those with by from as not no than then so if about after before into over per each
any have has had use used using""".split())

# Hand-made stand-in for what a neural model learns: words with similar meaning
# map to one shared concept. A real embedding model learns this from data.
SYNONYMS = {
    "hand": "due", "submit": "due", "deadline": "due", "turn": "due",
    "halfway": "milestone", "midway": "milestone", "mid": "milestone",
    "shortest": "quiet", "quietest": "quiet", "emptiest": "quiet", "least": "quiet",
    "queue": "wait", "line": "wait", "waiting": "wait", "crowded": "busy",
    "money": "credit", "funds": "credit", "dollars": "credit", "cash": "credit",
    "ask": "request", "appeal": "request", "grades": "grade", "scores": "grade",
    "group": "team", "squad": "team", "funding": "credit",
    "stuck": "error", "failed": "error", "crash": "error",
}


def stem(w):
    for suf in ("ing", "ed", "es", "s"):
        if len(w) > len(suf) + 2 and w.endswith(suf):
            return w[: -len(suf)]
    return w


def tokenize(text, drop_stop=True):
    words = re.findall(r"[a-z0-9]+", text.lower())
    return [stem(w) for w in words if not (drop_stop and w in STOPWORDS)]


_SYN = {stem(k): v for k, v in SYNONYMS.items()}   # tokens are stemmed, so stem the keys too


def concept(token):
    return _SYN.get(token, token)


def split_sentences(text):
    parts = re.split(r"(?<=[.!?])\s+", text.strip())
    return [p for p in parts if p]

# ---------------------------------------------------------------- chunking

@dataclass
class Chunk:
    id: str
    doc_id: str
    text: str
    context: str = ""          # optional prefix added by contextual retrieval

    @property
    def indexed_text(self):
        return (self.context + " " + self.text).strip()


def chunk_fixed(doc_id, text, size=30, overlap=0):
    """Split into windows of `size` words. `overlap` words repeat between neighbors."""
    if size < 1 or not 0 <= overlap < size:
        raise ValueError("need size >= 1 and 0 <= overlap < size")
    words, step, out, i = text.split(), size - overlap, [], 0
    while i < len(words):
        out.append(Chunk(f"{doc_id}#{len(out)}", doc_id, " ".join(words[i:i + size])))
        if i + size >= len(words):
            break
        i += step
    return out


def chunk_sentences(doc_id, text, max_words=30):
    """Pack whole sentences into chunks of at most max_words (a sentence is never cut)."""
    out, cur, n = [], [], 0
    for s in split_sentences(text):
        w = len(s.split())
        if cur and n + w > max_words:
            out.append(Chunk(f"{doc_id}#{len(out)}", doc_id, " ".join(cur)))
            cur, n = [], 0
        cur.append(s)
        n += w
    if cur:
        out.append(Chunk(f"{doc_id}#{len(out)}", doc_id, " ".join(cur)))
    return out


def add_context(chunks, doc_titles, summarize=None):
    """Contextual retrieval, toy version: prefix each chunk with where it came from.
    A real system asks an LLM to write the prefix. `summarize(chunk, title)` is that hook."""
    summarize = summarize or (lambda c, t: f"From the document '{t}'.")
    return [Chunk(c.id, c.doc_id, c.text, summarize(c, doc_titles[c.doc_id])) for c in chunks]

# ------------------------------------------------------------------- BM25

class BM25:
    """Okapi BM25 with the common non-negative IDF: ln(1 + (N - n + 0.5) / (n + 0.5))."""
    def __init__(self, chunks, k1=1.5, b=0.75):
        self.chunks, self.k1, self.b = chunks, k1, b
        self.docs = [tokenize(c.indexed_text) for c in chunks]
        self.N = len(self.docs)
        self.avgdl = sum(len(d) for d in self.docs) / max(self.N, 1)
        self.df = {}
        for d in self.docs:
            for t in set(d):
                self.df[t] = self.df.get(t, 0) + 1

    def idf(self, term):
        n = self.df.get(term, 0)
        return math.log(1 + (self.N - n + 0.5) / (n + 0.5))

    def score(self, query, i):
        d, s = self.docs[i], 0.0
        for t in set(tokenize(query)):
            f = d.count(t)
            if f:
                s += self.idf(t) * f * (self.k1 + 1) / (f + self.k1 * (1 - self.b + self.b * len(d) / self.avgdl))
        return s

    def search(self, query, k=5):
        scored = [(self.score(query, i), i) for i in range(self.N)]
        scored = [x for x in scored if x[0] > 0]
        scored.sort(key=lambda x: (-x[0], x[1]))
        return [(self.chunks[i].id, s) for s, i in scored[:k]]

# ------------------------------------------- toy dense retriever (stand-in)

def _bucket(feature, dim):
    h = zlib.crc32(feature.encode())
    return h % dim, (1.0 if (h >> 16) & 1 else -1.0)


def _normalize(v):
    n = math.sqrt(sum(x * x for x in v))
    return [x / n for x in v] if n else v


class ToyEmbedder:
    """Hashed bag of concepts + letter triples. NOT a neural model."""
    def __init__(self, dim=256):
        self.dim, self.idf = dim, {}

    def fit(self, texts):
        df, n = {}, len(texts)
        for t in texts:
            for w in {concept(x) for x in tokenize(t)}:
                df[w] = df.get(w, 0) + 1
        self.idf = {w: math.log(1 + n / c) for w, c in df.items()}
        return self

    def embed(self, text):
        v = [0.0] * self.dim
        for tok in tokenize(text):
            c = concept(tok)
            w = self.idf.get(c, math.log(1 + 10))
            for feat, weight in [("w:" + c, 1.0)] + [("t:" + c[i:i + 3], 0.15) for i in range(max(len(c) - 2, 0))]:
                j, sgn = _bucket(feat, self.dim)
                v[j] += sgn * weight * w
        return _normalize(v)

    def embed_tokens(self, text, dim=64):
        """One small vector per token, for late interaction."""
        small = ToyEmbedder(dim)
        small.idf = {}
        out = []
        for tok in tokenize(text):
            c, v = concept(tok), [0.0] * dim
            for feat, weight in [("w:" + c, 1.0)] + [("t:" + c[i:i + 3], 0.15) for i in range(max(len(c) - 2, 0))]:
                j, sgn = _bucket(feat, dim)
                v[j] += sgn * weight
            out.append(_normalize(v))
        return out


def cosine(a, b):
    return sum(x * y for x, y in zip(a, b))


class DenseIndex:
    def __init__(self, chunks, dim=256):
        self.chunks = chunks
        self.emb = ToyEmbedder(dim).fit([c.indexed_text for c in chunks])
        self.vecs = [self.emb.embed(c.indexed_text) for c in chunks]

    def search(self, query, k=5):
        q = self.emb.embed(query)
        scored = sorted(((cosine(q, v), i) for i, v in enumerate(self.vecs)), key=lambda x: (-x[0], x[1]))
        return [(self.chunks[i].id, s) for s, i in scored[:k] if s > 0]

# ------------------------------------------------------ fusion and rerank

def rrf(rankings, k=60, top=None):
    """Reciprocal Rank Fusion. `rankings` is a list of lists of ids, best first.
    score(d) = sum over lists of 1 / (k + rank), with rank starting at 1."""
    scores = {}
    for lst in rankings:
        for r, doc in enumerate(lst, start=1):
            scores[doc] = scores.get(doc, 0.0) + 1.0 / (k + r)
    fused = sorted(scores.items(), key=lambda x: (-x[1], x[0]))
    return fused[:top] if top else fused


def maxsim(query_vecs, doc_vecs):
    """Late interaction score: for each query token take its best match in the doc, then sum."""
    if not doc_vecs:
        return 0.0
    return sum(max(cosine(q, d) for d in doc_vecs) for q in query_vecs)


def rerank_late_interaction(query, chunks_by_id, ids, embedder):
    """Re-order candidate ids by MaxSim. Token vectors for chunks could be precomputed offline."""
    qv = embedder.embed_tokens(query)
    scored = [(maxsim(qv, embedder.embed_tokens(chunks_by_id[i].indexed_text)), i) for i in ids]
    scored.sort(key=lambda x: (-x[0], x[1]))
    return [(i, s) for s, i in scored]

# ---------------------------------------------------------------- metrics

def hit_at_k(ranked, relevant, k):
    return 1.0 if set(ranked[:k]) & set(relevant) else 0.0


def recall_at_k(ranked, relevant, k):
    return len(set(ranked[:k]) & set(relevant)) / len(relevant) if relevant else 0.0


def mrr(ranked, relevant, k=None):
    for r, d in enumerate(ranked[:k] if k else ranked, start=1):
        if d in relevant:
            return 1.0 / r
    return 0.0


def ndcg_at_k(ranked, relevant, k):
    """Binary relevance. DCG = sum rel_i / log2(i + 1)."""
    dcg = sum(1.0 / math.log2(i + 1) for i, d in enumerate(ranked[:k], start=1) if d in relevant)
    ideal = sum(1.0 / math.log2(i + 1) for i in range(1, min(len(relevant), k) + 1))
    return dcg / ideal if ideal else 0.0


def relevant_ids(chunks, needle):
    """Gold labels for the lab: a chunk is relevant if it contains the answer text."""
    return [c.id for c in chunks if needle.lower() in c.text.lower()]

# --------------------------------------------------------------- generator

NO_ANSWER = "I don't know based on the sources I have."


@dataclass
class Answer:
    text: str
    citations: list = field(default_factory=list)
    abstained: bool = False


def mock_llm_answer(question, retrieved, min_overlap=2):
    """A mock 'LLM'. It picks the source sentence that best matches the question.
    Matches on rarer concepts count more. It answers only if at least `min_overlap`
    concepts match, and it cites the chunk. Otherwise it abstains.
    `retrieved` is a list of Chunk objects, best first."""
    q = {concept(t) for t in tokenize(question)}
    pool = [(ch, s, {concept(t) for t in tokenize(s)}) for ch in retrieved for s in split_sentences(ch.text)]
    df = {}
    for _, _, cs in pool:
        for c in cs:
            df[c] = df.get(c, 0) + 1
    best, best_w, best_chunk = None, 0.0, None
    for ch, sent, cs in pool:
        hits = q & cs
        if len(hits) < min_overlap:
            continue
        w = sum(math.log(1 + len(pool) / df[c]) for c in hits)
        if w > best_w:
            best, best_w, best_chunk = sent, w, ch
    if best is None:
        return Answer(NO_ANSWER, [], True)
    return Answer(f"{best} [{best_chunk.id}]", [best_chunk.id], False)


def check_citations(answer, retrieved):
    """True when every cited id was really in the retrieved set (a basic faithfulness check)."""
    have = {c.id for c in retrieved}
    return all(c in have for c in answer.citations)

# --------------------------------------------------------------- agent loop

SMALL_TALK = {"hi", "hello", "hey", "thanks", "thank", "bye"}


def needs_retrieval(question):
    """Decision 1: retrieve or answer directly? Greetings need no sources."""
    toks = re.findall(r"[a-z]+", question.lower())
    return not (toks and len(toks) <= 3 and all(t in SMALL_TALK for t in toks))


def rewrite_query(question, retrieved):
    """Decision 2: what to ask next. Toy rule: (1) add the shared concept word for any
    question word that has one, which fixes vocabulary mismatch; (2) add the most
    frequent new content words from what we found. A real agent asks an LLM to rewrite."""
    toks = tokenize(question)
    extra = [concept(t) for t in toks if concept(t) != t]
    seen, counts = set(toks), {}
    for ch in retrieved[:2]:
        for t in tokenize(ch.text):
            if t not in seen and not t.isdigit():
                counts[t] = counts.get(t, 0) + 1
    extra += [t for t, _ in sorted(counts.items(), key=lambda x: (-x[1], x[0]))[:2]]
    return (question + " " + " ".join(extra)).strip()


def agentic_answer(question, search_fn, chunks_by_id, max_rounds=3, k=3):
    """Retrieve, check, maybe retry. Returns (Answer, trace).
    search_fn(query, k) -> list of ids. Decision 3: stop when grounded, or after max_rounds."""
    trace = []
    if not needs_retrieval(question):
        trace.append("skip retrieval: small talk")
        return Answer("Hello! Ask me about the course.", [], False), trace
    query, retrieved = question, []
    for rnd in range(1, max_rounds + 1):
        ids = search_fn(query, k * rnd)   # widen the search each round
        retrieved = [chunks_by_id[i] for i in ids]
        ans = mock_llm_answer(question, retrieved)
        trace.append(f"round {rnd}: query={query!r} -> {ids} -> {'abstain' if ans.abstained else 'grounded'}")
        if not ans.abstained:
            return ans, trace
        new_query = rewrite_query(question, retrieved)
        if new_query == query:
            trace.append("stop: the query did not change")
            return Answer(NO_ANSWER, [], True), trace
        query = new_query
    trace.append("stop: round limit reached")
    return Answer(NO_ANSWER, [], True), trace

# ------------------------------------------------------- the lab corpus
# A small, made-up help center for an imaginary course. All facts are invented.

DOCS = {
    "policies": ("Course Policies",
        "The project milestone is due in week 6 on Thursday at 11:59 pm. "
        "Late days cannot be used on the milestone. The final report is due in week 10. "
        "Each student has four late days to use on homework. "
        "Regrade requests must be sent within seven days of the grade release. "
        "Send every regrade request through the course forum, not by email."),
    "office": ("Office Hours",
        "Office hours run Monday to Friday. The Thursday 9 am session is the quietest, "
        "with about 3 students waiting on average. The Tuesday 4 pm session is the busiest, "
        "with about 14 students waiting. Bring your code and a short description of the bug."),
    "setup": ("Lab Setup Guide",
        "Install Python 3.11 and create a virtual environment. Install the course packages "
        "with the provided requirements file. Error ERR-4042 means the API key file is missing. "
        "Create a file named .env in the project folder and add your key. "
        "Error ERR-5100 means the request timed out. Wait one minute and retry."),
    "credits": ("Cloud Credits",
        "Each team gets 50 dollars of cloud credit for the quarter. "
        "Credit requests go through the staff form. Unused credit expires on the last day of the quarter."),
    "dining": ("Campus Dining",
        "The dining hall opens at 7 am. Late dinner is served until 9 pm on weekdays. "
        "The Thursday market sells snacks until noon. Report lost items to the front desk."),
}

# (question, answer text that must appear in a relevant chunk)
QUERIES = [
    ("When do I hand in the halfway project checkpoint?", "project milestone is due"),
    ("What does ERR-4042 mean?", "ERR-4042 means"),
    ("Which session has the shortest queue?", "Thursday 9 am session is the quietest"),
    ("How much cloud money does each team get?", "50 dollars"),
    ("How many days do I have to ask for a regrade?", "within seven days"),
    ("Can I use late days on the milestone?", "Late days cannot be used"),
]
UNANSWERABLE = "What is the name of the instructor's cat?"


def build_chunks(strategy="sentence", size=22, overlap=5, contextual=False):
    chunks = []
    for doc_id, (title, text) in DOCS.items():
        if strategy == "fixed":
            chunks += chunk_fixed(doc_id, text, size, 0)
        elif strategy == "overlap":
            chunks += chunk_fixed(doc_id, text, size, overlap)
        else:
            chunks += chunk_sentences(doc_id, text, size)
    if contextual:
        chunks = add_context(chunks, {k: v[0] for k, v in DOCS.items()})
    return chunks


class Pipeline:
    """Everything wired together: chunks -> BM25 + dense -> RRF -> optional rerank."""
    def __init__(self, chunks):
        self.chunks = chunks
        self.by_id = {c.id: c for c in chunks}
        self.bm25, self.dense = BM25(chunks), DenseIndex(chunks)

    def search(self, query, k=3, mode="hybrid", depth=10):
        if mode == "bm25":
            return [i for i, _ in self.bm25.search(query, k)]
        if mode == "dense":
            return [i for i, _ in self.dense.search(query, k)]
        fused = rrf([[i for i, _ in self.bm25.search(query, depth)],
                     [i for i, _ in self.dense.search(query, depth)]], top=depth)
        ids = [i for i, _ in fused]
        if mode == "hybrid+rerank":
            ids = [i for i, _ in rerank_late_interaction(query, self.by_id, ids, self.dense.emb)]
        return ids[:k]


def evaluate(pipe, mode, k=3):
    rows = []
    for q, needle in QUERIES:
        rel = relevant_ids(pipe.chunks, needle)
        ranked = pipe.search(q, k=10, mode=mode)
        rows.append((hit_at_k(ranked, rel, k), recall_at_k(ranked, rel, k), mrr(ranked, rel, 10), ndcg_at_k(ranked, rel, k)))
    n = len(rows)
    return [round(sum(r[i] for r in rows) / n, 2) for i in range(4)]


def _demo():
    print("== chunking ==")
    for strat in ("fixed", "overlap", "sentence"):
        cs = build_chunks(strat)
        print(f"{strat:9s}: {len(cs)} chunks; first of 'policies': {cs[0].text[:60]!r}...")
    chunks = build_chunks("sentence")
    pipe = Pipeline(chunks)
    print("\n== retrieval metrics over 6 questions (k=3): hit, recall, MRR, nDCG ==")
    for mode in ("bm25", "dense", "hybrid", "hybrid+rerank"):
        print(f"{mode:14s}", evaluate(pipe, mode))
    ctx_pipe = Pipeline(build_chunks("sentence", contextual=True))
    print(f"{'hybrid+context':14s}", evaluate(ctx_pipe, "hybrid"))
    print("\n== same questions, hybrid, three chunkers (size 12 words) ==")
    print("(a question scores 0 when no chunk contains the whole answer text)")
    for strat in ("fixed", "overlap", "sentence"):
        cs = build_chunks(strat, size=12, overlap=4)
        print(f"{strat:9s} chunks={len(cs):2d}", evaluate(Pipeline(cs), "hybrid"))

    print("\n== rank of the first relevant chunk, per question (- = not in top 10) ==")
    print(f"{'mode':14s}" + " ".join(f"Q{i + 1:<2d}" for i in range(len(QUERIES))))
    for mode in ("bm25", "dense", "hybrid", "hybrid+rerank"):
        cells = []
        for q, needle in QUERIES:
            rel = relevant_ids(pipe.chunks, needle)
            ranked = pipe.search(q, 10, mode)
            cells.append(str(next((r for r, d in enumerate(ranked, 1) if d in rel), "-")))
        print(f"{mode:14s}" + " ".join(f"{c:<3s}" for c in cells))

    print("\n== grounded answers (hybrid + rerank) ==")
    for q, _ in QUERIES[:3] + [(UNANSWERABLE, "")]:
        ids = pipe.search(q, 3, mode="hybrid+rerank")
        a = mock_llm_answer(q, [pipe.by_id[i] for i in ids])
        print("Q:", q, "\nA:", a.text)

    print("\n== agent loop ==")
    for q in ("hello", "What does ERR-4042 mean?", UNANSWERABLE):
        a, trace = agentic_answer(q, lambda s, k: pipe.search(s, k), pipe.by_id)
        print("Q:", q)
        for line in trace:
            print("   ", line)
        print("   ->", a.text)
    q = "What is the cash for each group?"
    a, trace = agentic_answer(q, lambda s, k: [i for i, _ in pipe.bm25.search(s, k)], pipe.by_id, k=1)
    print("Q (BM25 only, so the first query misses):", q)
    for line in trace:
        print("   ", line)
    print("   ->", a.text)


if __name__ == "__main__":
    _demo()