import torch
import torch.utils.benchmark as benchmark

from torch import Tensor
from torchvision.transforms.functional_tensor import _assert_image_tensor, _assert_channels, get_dimensions, rgb_to_grayscale, _blurred_degenerate_image
from torchvision.transforms import functional_tensor as V1



def _blend(img1: Tensor, img2: Tensor, ratio: float) -> Tensor:
    if ratio == 1.0:
        return img1
    ratio = float(ratio)
    bound = 1.0 if img1.is_floating_point() else 255.0

    if img2.is_floating_point():
        # Since img2 is float, we can do in-place ops on it. It's a throw-away tensor.
        # Our strategy is to convert img1 to float and copy it to avoid in-place modifications,
        # update img2 in-place and add it on the result with an in-place op.
        result = img1 * ratio
        img2.mul_(1.0 - ratio)
        result.add_(img2)
    else:
        # Since img2 is not float, we can't do in-place ops on it.
        # To minimize copies/adds/muls we first convert img1 to float by multiplying it with ratio/(1-ratio).
        # This permits us to add img2 in-place to it, without further copies.
        # To ensure we have the correct result at the end, we multiply in-place with (1-ratio).
        result = img1 * (ratio / (1.0 - ratio))
        result.add_(img2).mul_(1.0 - ratio)

    return result.clamp_(0, bound).to(img1.dtype)


############ COPY PASTE FROM functional_tensor.py
def adjust_brightness(img: Tensor, brightness_factor: float) -> Tensor:
    if brightness_factor < 0:
        raise ValueError(f"brightness_factor ({brightness_factor}) is not non-negative.")

    _assert_image_tensor(img)

    _assert_channels(img, [1, 3])

    return _blend(img, torch.zeros_like(img), brightness_factor)


def adjust_contrast(img: Tensor, contrast_factor: float) -> Tensor:
    if contrast_factor < 0:
        raise ValueError(f"contrast_factor ({contrast_factor}) is not non-negative.")

    _assert_image_tensor(img)

    _assert_channels(img, [3, 1])
    c = get_dimensions(img)[0]
    dtype = img.dtype if torch.is_floating_point(img) else torch.float32
    if c == 3:
        mean = torch.mean(rgb_to_grayscale(img).to(dtype), dim=(-3, -2, -1), keepdim=True)
    else:
        mean = torch.mean(img.to(dtype), dim=(-3, -2, -1), keepdim=True)

    return _blend(img, mean, contrast_factor)


def adjust_saturation(img: Tensor, saturation_factor: float) -> Tensor:
    if saturation_factor < 0:
        raise ValueError(f"saturation_factor ({saturation_factor}) is not non-negative.")

    _assert_image_tensor(img)

    _assert_channels(img, [1, 3])

    if get_dimensions(img)[0] == 1:  # Match PIL behaviour
        return img

    return _blend(img, rgb_to_grayscale(img), saturation_factor)


def adjust_sharpness(img: Tensor, sharpness_factor: float) -> Tensor:
    if sharpness_factor < 0:
        raise ValueError(f"sharpness_factor ({sharpness_factor}) is not non-negative.")

    _assert_image_tensor(img)

    _assert_channels(img, [1, 3])

    if img.size(-1) <= 2 or img.size(-2) <= 2:
        return img

    return _blend(img, _blurred_degenerate_image(img), sharpness_factor)


############### BENCHMARKS

def bench_torch(fn, data, factor, runtime=10):
    for _ in range(10):
        fn(data, factor)
    results = benchmark.Timer(
                stmt=f"fn(data, {factor})",
                globals={
                    "data": data,
                    "fn": fn,
                },
                num_threads=torch.get_num_threads(),
            ).blocked_autorange(min_run_time=runtime)
    return results.median


data = torch.randint(0, 256, (10, 3, 128, 128))
data_f = data / 255.0
factor = 1.5
bench_fn = bench_torch


devices = ["cpu"]
if torch.cuda.is_available():
    devices.append("cuda")
for fn_name in ["adjust_brightness", "adjust_contrast", "adjust_saturation", "adjust_sharpness"]:
    fn_v1 = V1.__dict__[fn_name]
    fn_v2 = locals()[fn_name]
    for device in devices:
        for dtype, img in [(torch.uint8, data.to(device=device)), (torch.float32, data_f.to(device=device))]:
            v1_time = bench_fn(fn_v1, img, factor)
            v2_time = bench_fn(fn_v2, img, factor)

            winner = "v1" if v1_time < v2_time else "v2"
            print(f"{fn_name} - Winner: {winner}(device={device}, dtype={dtype}) - v1: {v1_time}, v2: {v2_time} - Diff: {100*(v2_time-v1_time)/v1_time:.2f}%")
            try:
                torch.testing.assert_close(fn_v1(img, factor), fn_v2(img, factor))
            except Exception:
                print("WARNING: fv1(x) != fv2(x)")