"""How much room does everything take on the disk?

    python sizes.py

Run it after "python train.py lora" and "python train.py full", with the same
Python. It adds up the model you downloaded, the two folders that training
made, and the Python libraries.
"""
import sys
from pathlib import Path

from huggingface_hub import snapshot_download

MODEL = "HuggingFaceTB/SmolLM2-135M-Instruct"
PAGE_FILES = Path(__file__).resolve().parent.parent


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):
    if not Path(folder).exists():
        print(f"{label:34} is not there yet")
        return 0.0
    size, files = megabytes(folder)
    print(f"{label:34}{size:10.1f} MB in {files:,} files")
    return size


def main():
    downloaded = snapshot_download(MODEL, local_files_only=True)
    show("the model, as downloaded", downloaded)
    lora = show("tuned_lora (the add-ons)", PAGE_FILES / "tuned_lora")
    full = show("tuned_full (a whole model)", PAGE_FILES / "tuned_full")
    if lora and full:
        print(f"tuned_full is {full / lora:,.0f} times as big as tuned_lora.")
    print()
    for file in sorted((PAGE_FILES / "tuned_lora").glob("*")):
        print(f"    tuned_lora/{file.name:28}{file.stat().st_size:14,} bytes")
    for file in sorted((PAGE_FILES / "tuned_full").glob("*")):
        print(f"    tuned_full/{file.name:28}{file.stat().st_size:14,} bytes")
    print()
    show("the Python libraries (the venv)", sys.prefix)
    for name in ["train.jsonl", "test_generated.jsonl", "test_handwritten.jsonl"]:
        print(f"{name:34}{(PAGE_FILES / name).stat().st_size / 1e3:10.1f} KB")


if __name__ == "__main__":
    main()