"""How big is everything, and how long is the handbook?

    python sizes.py

Run it with the page's Python, after the page's files have downloaded the two
models. It adds up the models as downloaded and the Python libraries, counts
the numbers in each model, and counts the handbook: its passages, its words,
and its tokens as the reader model counts them.
"""
import sys
from pathlib import Path

import torch
from huggingface_hub import snapshot_download
from sentence_transformers import SentenceTransformer
from transformers import AutoConfig, 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)


def megabytes(folder):
    """The size of every file in a folder, in MB, and how many files there are."""
    files = [file for file in Path(folder).rglob("*") if file.is_file()]
    return sum(file.stat().st_size for file in files) / 1e6, len(files)


def show(label, folder):
    size, files = megabytes(folder)
    print(f"{label:38}{size:9.1f} MB in {files:,} files")


def count_numbers(model):
    return sum(tensor.numel() for tensor in model.parameters())


def main():
    show("the embedding model, as downloaded", snapshot_download(answer.EMBEDDER, local_files_only=True))
    show("the reader model, as downloaded", snapshot_download(answer.MODEL, local_files_only=True))
    show("the Python libraries (the venv)", sys.prefix)
    print()

    embedder = SentenceTransformer(answer.EMBEDDER)
    print(f"the embedding model has {count_numbers(embedder):,} numbers,"
          f" and reads the first {embedder.max_seq_length} tokens of a text")
    reader = AutoModelForCausalLM.from_pretrained(answer.MODEL, dtype=torch.float32)
    limit = AutoConfig.from_pretrained(answer.MODEL).max_position_embeddings
    print(f"the reader model has {count_numbers(reader):,} numbers,"
          f" and can read {limit:,} tokens at once")
    print()

    passages = answer.read_passages(PAGE_FILES / "handbook.txt")
    text = (PAGE_FILES / "handbook.txt").read_text(encoding="utf-8")
    tokenizer = AutoTokenizer.from_pretrained(answer.MODEL)
    tokens = len(tokenizer(text, add_special_tokens=False)["input_ids"])
    words = [len(passage.split()) for passage in passages]
    print(f"handbook.txt: {len(passages)} passages, {len(text.split())} words,"
          f" {tokens:,} tokens, {len(text.encode('utf-8')) / 1e3:.1f} KB")
    print(f"the shortest passage has {min(words)} words and the longest {max(words)}")
    longest = max(len(embedder.tokenizer(passage)["input_ids"]) for passage in passages)
    print(f"the longest passage is {longest} tokens for the embedding model")
    questions = answer.read_questions(PAGE_FILES / "questions.jsonl")
    with_answer = sum(question["answer_in"] is not None for question in questions)
    print(f"questions.jsonl: {len(questions)} questions, {with_answer} that the handbook"
          f" answers and {len(questions) - with_answer} that it does not")


if __name__ == "__main__":
    main()