import itertools
import pathlib
import pickle

import torch
from torch.utils import benchmark
import functools

from torchvision.prototype.transforms import functional as F

description = "PR"  # "main", "PR"


def make_inputs(*, input_dtype, output_dtype, device, shape=(3, 512, 512)):
    if input_dtype.is_floating_point:
        image = torch.rand(shape, dtype=input_dtype, device=device)
    else:
        image = torch.randint(0, torch.iinfo(input_dtype).max + 1, shape, dtype=input_dtype, device=device)
    return image, output_dtype


sub_labels_and_input_fns = [
    ("float to float", functools.partial(make_inputs, input_dtype=torch.float32, output_dtype=torch.float64)),
    ("float to   int", functools.partial(make_inputs, input_dtype=torch.float32, output_dtype=torch.uint8)),
    ("  int to float", functools.partial(make_inputs, input_dtype=torch.uint8, output_dtype=torch.float32)),
    ("  int to   int (down)", functools.partial(make_inputs, input_dtype=torch.int32, output_dtype=torch.uint8)),
    ("  int to   int (up)", functools.partial(make_inputs, input_dtype=torch.uint8, output_dtype=torch.int32)),
]


timers = [
    benchmark.Timer(
        stmt="convert_image_dtype(*inputs)",
        globals=dict(
            convert_image_dtype=F.convert_image_dtype,
            inputs=inputs_fn(device=device),
        ),
        label="convert_image_dtype perf improvements",
        sub_label=f"{device:4} / {sub_label}",
        description=description,
        num_threads=num_threads,
    )
    for (sub_label, inputs_fn), device in itertools.product(sub_labels_and_input_fns, ["cpu", "cuda"])
    for num_threads in ([1, 2, 4] if device == "cpu" else [1])
]

measurements = [timer.blocked_autorange(min_run_time=5) for timer in timers]


with open(f"{description}.measurements", "wb") as fh:
    pickle.dump(measurements, fh)

measurements = []
for file in pathlib.Path(".").glob("*.measurements"):
    with open(file, "rb") as fh:
        measurements.extend(pickle.load(fh))

comparison = benchmark.Compare(measurements)
comparison.trim_significant_figures()
comparison.print()