"""What does the decider-ai library need to train a model?

    python decider_training.py

Run it with a Python that has the library:  pip install "decider-ai[train]"
It trains nothing. It prints, straight from the library:
  - the lines of its training script that decide what a training run needs;
  - what its README says about its models and about training;
  - the first lines of scripts/train.sh, which is in the library's repository
    on GitHub and not in what pip installs, so it is read from GitHub;
  - the options of the training script;
  - whether this computer has the kind of graphics card the script asks for.
"""
import subprocess
import sys
import urllib.request
from importlib import metadata, util
from pathlib import Path

import torch

WANTED = {
    "train.py": ['"--model"', '"--lr"', '"--max_tokens"', ".cuda()", "torch.optim"],
    "model.py": ["def __init__", "gradient_checkpointing"],
}
README_LINES = ["| decider-2b **v11** |", "| decider-12b **v2** |",
                "merged LoRA (rank 32) trained on"]
TRAIN_SH = "https://raw.githubusercontent.com/Mapika/decider/main/scripts/train.sh"


def show_lines(package, name):
    """Print the lines of one library file that hold any of the wanted words."""
    lines = (package / name).read_text(encoding="utf-8").splitlines()
    print(f"decider/{name}, {len(lines)} lines:")
    for number, line in enumerate(lines, start=1):
        if any(words in line for words in WANTED[name]):
            print(f"  line {number:3}:  {line.strip()[:104]}")


def readme_section(readme, heading):
    """One section of the README, from its heading to the next heading."""
    start = readme.index(heading)
    end = readme.index("\n## ", start + len(heading))
    return readme[start:end].strip()


def show_train_sh():
    """The comment lines at the top of scripts/train.sh, and its learning rates."""
    try:
        request = urllib.request.Request(TRAIN_SH, headers={"User-Agent": "Mozilla/5.0"})
        text = urllib.request.urlopen(request, timeout=30).read().decode("utf-8")
    except OSError as problem:
        print(f"  could not read it: {problem}")
        return
    for number, line in enumerate(text.splitlines(), start=1):
        if line.startswith("# ") or line.startswith("LR="):
            print(f"  line {number:2}:  {line[:150]}")


def main():
    package = Path(util.find_spec("decider").origin).parent
    readme = metadata.metadata("decider-ai").get_payload()
    print(f"decider-ai {metadata.version('decider-ai')}, torch {torch.__version__}")
    print()
    for name in WANTED:
        show_lines(package, name)
    print()
    print("From the library's README, about its models:")
    for line in readme.splitlines():
        if any(words in line for words in README_LINES):
            print(" ", line.strip()[:150])
    print()
    print("From the library's README, about training:")
    print(readme_section(readme, "## Train your own"))
    print()
    print(f"scripts/train.sh, read from {TRAIN_SH}")
    show_train_sh()
    print()
    print("python -m decider.train --help")
    printed = subprocess.run([sys.executable, "-m", "decider.train", "--help"],
                             capture_output=True, text=True, encoding="utf-8")
    print(printed.stdout.strip() or printed.stderr.strip()[-800:])
    print()
    if torch.cuda.is_available():
        print(f"This computer has a CUDA graphics card: {torch.cuda.get_device_name(0)}.")
    else:
        print("This computer: torch.cuda.is_available() is False.")
        print("train.py moves the model to a CUDA graphics card with .cuda(),"
              " so it cannot train here.")


if __name__ == "__main__":
    main()