"""Manual retrieval-augmented generation for a document Q&A lesson.

Both the embedding model and the generation model run on CPU. The generation
model is loaded only when a user asks for an actual generated answer.
"""

import numpy as np
import torch
from datasets import load_dataset
from sentence_transformers import SentenceTransformer
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
GENERATION_MODEL = "google/flan-t5-small"
MIN_SCORE = 0.30  # Teaching heuristic, NOT calibrated confidence.
NO_ANSWER = "I cannot answer this from the supplied passages."
DATASET_ID = "rajpurkar/squad"
DATASET_REVISION = "7b6d24c440a36b6815f21b70d25016731768db1f"

corpus = []
source_questions = []
embedder = None
passage_embeddings = None
tokenizer = None
generator = None


def load_passage_sample(embedding_tokenizer):
    """Stream a reproducible sample of 30 distinct SQuAD contexts."""
    rows = load_dataset(
        DATASET_ID,
        split="train",
        revision=DATASET_REVISION,
        streaming=True,
    )
    passages, questions = [], []
    seen_contexts = set()
    title_counts = {}

    for row in rows.take(6000):
        title = row["title"].replace("_", " ")
        context = row["context"].strip()
        if (
            context in seen_contexts
            or title_counts.get(title, 0) >= 3
            or len(title) > 35
            or title == "Genocide"
        ):
            continue

        token_count = len(
            embedding_tokenizer(
                title + ". " + context,
                truncation=False,
                verbose=False,
            )["input_ids"]
        )
        if token_count > 200:
            continue

        passage_id = f"P{len(passages) + 1:02d}"
        passages.append(
            {
                "id": passage_id,
                "title": title,
                "text": context,
                "source": "SQuAD train / Wikipedia: " + title,
                "squad_id": row["id"],
            }
        )
        questions.append(
            {
                "query": row["question"],
                "expected_ids": [passage_id],
                "squad_id": row["id"],
                "answer_spans": row["answers"]["text"],
            }
        )
        seen_contexts.add(context)
        title_counts[title] = title_counts.get(title, 0) + 1
        if len(passages) == 30:
            break

    if len(passages) != 30:
        raise RuntimeError(
            "The Hub sample did not produce 30 passages. "
            "Check dataset access and the specified revision."
        )
    return passages, questions


def make_evaluation_cases(questions):
    """Create answerable and unanswerable evaluation examples."""
    if len(questions) != 30:
        raise ValueError("Load the Hub sample before creating evaluation cases.")