"""Run decide.py or agent.py for real, with the model on this computer, and time it.

    python model_timing.py decide.py
    python model_timing.py agent.py

It needs what those files need: the decider-ai library (pip install decider-ai) and the
model, which the first run downloads. The file itself is not changed. Before it runs, the
library's Decider is swapped for one that writes down how long loading and each question
take.

What the file prints goes to the normal output, so this also makes its output file:

    python model_timing.py agent.py > ../agent_output.txt

The times, the computer and one whole answer go to the error stream.
"""
import json
import pathlib
import platform
import runpy
import sys
import time

ASSETS = pathlib.Path(__file__).resolve().parent.parent      # the tutorial's files: one folder up from checks/


def main():
    if len(sys.argv) != 2 or sys.argv[1] not in ("decide.py", "agent.py"):
        print(__doc__)
        return 2
    started = time.perf_counter()
    import torch
    import decider.infer
    from importlib import metadata
    imported = time.perf_counter() - started
    print(f"python {platform.python_version()} on {platform.system()} {platform.machine()}, "
          f"decider-ai {metadata.version('decider-ai')}, torch {torch.__version__}, "
          f"transformers {metadata.version('transformers')}, cuda {torch.cuda.is_available()}, "
          f"threads {torch.get_num_threads()}", file=sys.stderr)
    real = decider.infer.Decider
    times = {"load": 0.0, "questions": []}
    kept = {}

    class TimedDecider(real):
        def __init__(self, *arguments, **options):
            begun = time.perf_counter()
            super().__init__(*arguments, **options)
            times["load"] = time.perf_counter() - begun
            kept["model"] = self
            print(f"the model is on: {self.dev}, as {self.name}", file=sys.stderr)

        def system_one(self, *arguments, **options):
            begun = time.perf_counter()
            answer = super().system_one(*arguments, **options)
            times["questions"].append(time.perf_counter() - begun)
            kept["last"] = answer
            return answer

    decider.infer.Decider = TimedDecider
    begun = time.perf_counter()
    runpy.run_path(str(ASSETS / sys.argv[1]), run_name="__main__")
    in_all = time.perf_counter() - begun + imported
    print(f"import {imported:.1f} s, load {times['load']:.1f} s, questions "
          f"{' '.join(f'{seconds:.2f}' for seconds in times['questions'])} s, "
          f"{in_all:.1f} s in all", file=sys.stderr)
    print("the last whole answer:", json.dumps(kept.get("last")), file=sys.stderr)
    return 0


if __name__ == "__main__":
    sys.exit(main())