from __future__ import annotations
import torch.utils.benchmark as benchmark
import torch
import cProfile
from monai.metrics.hausdorff_distance import compute_hausdorff_distance
from monai.metrics import utils
from pathlib import Path
import numpy as np

shapes = [
    #   (250, 250, 250),
    #   (100, 100, 100),
    (50, 50, 50),
]


def create_batch(shape, device, use_random=False, random_ratio=0.2):
    if not use_random:
        [seg_1, seg_2, _] = [
            create_spherical_seg_3d(radius=20, centre=(20, 20, 20), im_shape=shape),
            create_spherical_seg_3d(radius=20, centre=(19, 19, 19), im_shape=shape),
            None,
        ]
    else:
        [seg_1, seg_2, _] = [
            create_random_seg(random_ratio, shape),
            create_random_seg(random_ratio, shape),
            None,
        ]

    batch, n_class = 2, 3
    batch_seg_1 = (
        torch.tensor(seg_1)
        .unsqueeze(0)
        .unsqueeze(0)
        .repeat([batch, n_class, 1, 1, 1])
        .to(device)
    )
    batch_seg_2 = (
        torch.tensor(seg_2)
        .unsqueeze(0)
        .unsqueeze(0)
        .repeat([batch, n_class, 1, 1, 1])
        .to(device)
    )
    return batch_seg_1, batch_seg_2


def benchmark_compare():
    results = []
    ratios = [0.2]
    for shape in shapes:
        for device in ["cpu", "cuda"]:
            utils.has_cucim_distance_transform_edt = device != "cpu"
            for use_random in [True, False]:
                for ratio in ratios:
                    if not use_random and ratio != ratios[0]:
                        continue
                    batch_seg_1, batch_seg_2 = create_batch(
                        shape, device, use_random=use_random, random_ratio=ratio
                    )
                    t = benchmark.Timer(
                        stmt="compute_hausdorff_distance(y_pred=seg_1,y=seg_2)",
                        globals={
                            "seg_1": batch_seg_1,
                            "seg_2": batch_seg_2,
                            "compute_hausdorff_distance": compute_hausdorff_distance,
                        },
                        label=str(shape),
                        sub_label=f"random()>{ratio}"
                        if use_random
                        else "create_spherical_seg_3d",
                        description=device,
                        num_threads=1,
                    )
                    results.append(t.blocked_autorange(min_run_time=10))

    compare = benchmark.Compare(results)
    compare.print()


def profile():
    Path("./perf").mkdir(exist_ok=True)
    for shape in shapes:
        for device in ["cpu", "cuda"]:
            utils.has_cucim_distance_transform_edt = device != "cpu"
            for use_random in [True, False]:
                print(f"{shape} {device} {use_random}")
                batch_seg_1, batch_seg_2 = create_batch(
                    shape, device, use_random=use_random
                )
                compute_hausdorff_distance(y_pred=batch_seg_1, y=batch_seg_2)
                with cProfile.Profile() as pr:
                    compute_hausdorff_distance(y_pred=batch_seg_1, y=batch_seg_2)
                    pr.dump_stats(
                        f"./perf/hausdorff_{shape[0]}_{use_random}_{device}.prof"
                    )


def create_spherical_seg_3d(
    radius: float = 20.0,
    centre: tuple[int, int, int] = (49, 49, 49),
    im_shape: tuple[int, int, int] = (99, 99, 99),
    im_spacing: tuple[float, float, float] = (1.0, 1.0, 1.0),
) -> np.ndarray:
    """
    Return a 3D image with a sphere inside. Voxel values will be
    1 inside the sphere, and 0 elsewhere.

    Args:
        radius: radius of sphere (in terms of number of voxels, can be partial)
        centre: location of sphere centre.
        im_shape: shape of image to create.
        im_spacing: spacing of image to create.

    See also:
        :py:meth:`~create_test_image_3d`
    """
    # Create image
    image = np.zeros(im_shape, dtype=np.int32)
    spy, spx, spz = np.ogrid[: im_shape[0], : im_shape[1], : im_shape[2]]
    spy = spy.astype(float) * im_spacing[0]
    spx = spx.astype(float) * im_spacing[1]
    spz = spz.astype(float) * im_spacing[2]

    spy -= centre[0]
    spx -= centre[1]
    spz -= centre[2]

    circle = (spx * spx + spy * spy + spz * spz) <= radius * radius

    image[circle] = 1
    image[~circle] = 0
    return image


def create_random_seg(ratio, shape):
    seg = np.zeros(shape, dtype=bool)
    seg[np.random.random(shape) > ratio] = 1
    return seg


if __name__ == "__main__":
    benchmark_compare()
    profile()