from contextlib import ContextDecorator
from time import perf_counter

import monai
from monai.transforms import LoadImage

loader_cucim = LoadImage(
    reader=monai.data.wsi_reader.WSIReader, backend="cucim", image_only=True)


class Timer(ContextDecorator):
    def __init__(self, message):
        self.message = message
        self.end = None

    def elapsed_time(self):
        self.end = perf_counter()
        return self.end - self.start

    def __enter__(self):
        self.start = perf_counter()
        return self

    def __exit__(self, exc_type, exc, exc_tb):
        if not self.end:
            self.elapsed_time()
        print("{} : {}".format(self.message, self.end - self.start))


image_file = "/home/gbae/repo/cucim/notebooks/input/006388_0.tif"

with Timer("  Thread elapsed time (cucim loader)") as timer:
    image = loader_cucim(image_file)
print(f"    type: {type(image)}, shape: {image.shape}, dtype: {image.dtype}")