# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_cli.ipynb.

# %% auto 0
__all__ = ['cli']

# %% ../nbs/01_cli.ipynb 3
import platform

import torch
import timm
import wandb

from fastcore.script import *

from .benchmark import *

# %% ../nbs/01_cli.ipynb 4
@call_parse
def cli(wnb: str ="disabled",       # W&B mode. Accepted values: online, offline, disabled.
        wnb_run: str =None,    # W&B run name (auto-generate if None)
        wnb_project: str =None,
        wnb_entity: str =None,
        run_number = 1,             # A unique number to keep track over repeat runs
        model: str ="resnet50",     # TIMM Model name
        bs: int =32,                #Batch size
        size: int =224,             # (fake) image size
        fp16: bool =False,
        n_batches = 0,              # Run for N batches. Mututally exclusive with `n_seconds`
        n_seconds = 0,              # Run for N seconds. Mutually exclusive with `n_batches`
    ):

    if not n_batches and not n_seconds:
        print("Either `n_batches` or `n_seconds` must be non-zero")
        exit(1)
    if n_batches and n_seconds:
        print(f"`n_batches` and `n_seconds` are mutually explusive. {n_batches=}, {n_seconds=}")
        exit(1)

    assert wnb in ["online", "offline", "disabled"]
    
    stats = {
        "device_name" : torch.cuda.get_device_name(),
        "device_capability": str(torch.cuda.get_device_capability()),
        "pytorch": torch.version.__version__,
        "cuda": torch.version.cuda,
        "platform": platform.platform(),
        "model": model,
        "fp16": fp16,
        "bs": bs,
        "n_seconds": n_seconds if n_seconds else None,
        "n_batches": n_batches if n_batches else None,
        "run_number": run_number, 
        "gpu_mem" : torch.cuda.get_device_properties(0).total_memory,
    }

    print("\n".join([ f"{k}: {v}" for k, v in stats.items() ]))

    run =  wandb.init(mode=wnb, project=wnb_project, entity=wnb_entity, name=wnb_run)
    run.log(stats, step=0, commit=True)

    model = timm.create_model(model, pretrained=False)
    duration, n_items = benchmark(model, bs=bs, size=size, fp16=fp16, n_batches=n_batches, n_seconds=n_seconds)
    summary = {
        "duration": duration,
        "n_items": n_items,
        "throughput": n_items / duration,
    }

    print("\n".join([ f"{k}: {v}" for k, v in summary.items() ]))

    run.log(summary, step=1, commit=True)
    run.finish()



