"""Which small model reads the found passages best?

    python readers.py

For each model in READERS it answers the practice questions in
practice_questions.jsonl, each with the three passages the search found. It
uses answer.py's own functions, so a model is asked, and its reply is marked,
exactly as "python answer.py found" does it. The practice questions are not the
test questions: they exist so that the reader could be chosen without looking
at the test.

A chat template either has a place for documents or it has none. A model with
such a place is asked through it, as on the page. Every model is also asked the
other way: with the passages written into the instruction by hand.

The first run downloads every model in READERS, about 6 GB in all.
"""
import sys
import time
from pathlib import Path

import torch
from sentence_transformers import SentenceTransformer
from transformers import AutoModelForCausalLM, AutoTokenizer

PAGE_FILES = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(PAGE_FILES))
import answer  # noqa: E402  (the page's own file)

READERS = [
    "HuggingFaceTB/SmolLM2-135M-Instruct",      # the model of the fine-tuning page
    "ibm-granite/granite-4.0-350m",             # the reader on this page
    "HuggingFaceTB/SmolLM2-360M-Instruct",
    "Qwen/Qwen2.5-0.5B-Instruct",
    "ibm-granite/granite-4.0-1b",               # the bigger reader of bigger_reader.py
]
BY_HAND = (" Use only the handbook below. If the handbook does not say, answer:"
           " I don't know.\n\nHANDBOOK\n")


def tokens_by_hand(tokenizer, question_text, passages):
    """The prompt with the passages pasted into the instruction, and no documents."""
    chat = [
        {"role": "system", "content": answer.INSTRUCTION + BY_HAND + "\n\n".join(passages)},
        {"role": "user", "content": question_text},
    ]
    return tokenizer.apply_chat_template(
        chat, add_generation_prompt=True, return_tensors="pt")["input_ids"]


def tokens_by_template(tokenizer, question_text, passages):
    """The prompt as answer.py makes it: the passages go in as documents."""
    documents = [answer.as_document(passage) for passage in passages]
    return answer.prompt_as_tokens(tokenizer, question_text, documents)


def takes_documents(tokenizer, question_text, passages):
    """True when a prompt made with documents= really has the passages in it."""
    prompt_tokens = tokens_by_template(tokenizer, question_text, passages)
    prompt_text = tokenizer.decode(prompt_tokens[0])
    return answer.as_document(passages[0])["text"][:40] in prompt_text


def score(model, tokenizer, questions, found, as_tokens):
    """Right answers, right refusals, seconds a reply, and the replies that missed."""
    right = refused = 0
    missed = []
    started = time.perf_counter()
    for question, passages in zip(questions, found):
        prompt_tokens = as_tokens(tokenizer, question["question"], passages)
        reply = answer.reply_to(model, tokenizer, prompt_tokens)
        if not answer.is_right(reply, question):
            missed.append((question["question"], reply))
        elif question["answer_in"] is None:
            refused += 1
        else:
            right += 1
    return right, refused, (time.perf_counter() - started) / len(questions), missed


def main():
    passages = answer.read_passages(PAGE_FILES / "handbook.txt")
    questions = answer.read_questions(Path(__file__).with_name("practice_questions.jsonl"))
    embedder = SentenceTransformer(answer.EMBEDDER)
    passage_vectors = embedder.encode(passages, normalize_embeddings=True)
    found = [answer.nearest(embedder, passages, passage_vectors, question["question"])
             for question in questions]
    in_handbook = sum(question["answer_in"] is not None for question in questions)
    none = len(questions) - in_handbook
    print(f"{len(questions)} practice questions: the handbook answers {in_handbook}"
          f" and has nothing on {none}.")
    print()
    table = []
    for name in READERS:
        tokenizer = AutoTokenizer.from_pretrained(name)
        model = AutoModelForCausalLM.from_pretrained(name, dtype=torch.float32)
        numbers = sum(tensor.numel() for tensor in model.parameters())
        ways = [("by hand", tokens_by_hand)]
        if takes_documents(tokenizer, questions[0]["question"], found[0]):
            ways.append(("documents=", tokens_by_template))
            print(f"{name}: a prompt made with documents= has the passages in it")
        else:
            print(f"{name}: a prompt made with documents= has NO passage in it,")
            print("    and no error either. Its chat template has no place for documents.")
        print()
        for way, as_tokens in ways:
            right, refused, seconds, missed = score(model, tokenizer, questions, found, as_tokens)
            table.append(f"{name:37}{numbers / 1e6:6.0f} M   {way:11}"
                         f"{right:3} of {in_handbook}   {refused} of {none}   {seconds:5.1f} s")
            print(f"{name}, passages given {way}")
            for question_text, reply in missed:
                print(f"    {question_text}")
                print(f"        {reply}")
            print()
        del model
    print(f"{'model':37}{'numbers':>8}   {'passages':11}{'right':>9}   refused   a reply")
    for line in table:
        print(line)


if __name__ == "__main__":
    main()