# Copyright 2020 - 2021 MONAI Consortium
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import unittest
from typing import Tuple

import numpy as np
import torch
from ignite.engine import Engine

from monai.handlers import SurfaceDistance


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)
) -> 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

    See also:
        :py:meth:`~create_test_image_3d`
    """
    # Create image
    image = np.zeros(im_shape, dtype=np.int32)
    spy, spx, spz = np.ogrid[
        -centre[0] : im_shape[0] - centre[0], -centre[1] : im_shape[1] - centre[1], -centre[2] : im_shape[2] - centre[2]
    ]
    circle = (spx * spx + spy * spy + spz * spz) <= radius * radius

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


sampler_sphere = torch.Tensor(create_spherical_seg_3d(radius=20, centre=(20, 20, 20))).unsqueeze(0).unsqueeze(0)
# test input a list of channel-first tensor
sampler_sphere_gt = [torch.Tensor(create_spherical_seg_3d(radius=20, centre=(10, 20, 20))).unsqueeze(0)]
sampler_sphere_zeros = torch.zeros_like(sampler_sphere)

TEST_SAMPLE_1 = [sampler_sphere, sampler_sphere_gt]
TEST_SAMPLE_2 = [sampler_sphere_gt, sampler_sphere_gt]
TEST_SAMPLE_3 = [sampler_sphere_zeros, sampler_sphere_gt]
TEST_SAMPLE_4 = [sampler_sphere_zeros, sampler_sphere_zeros]


class TestHandlerSurfaceDistance(unittest.TestCase):
    # TODO test multi node Surface Distance

    def test_compute(self):
        sur_metric = SurfaceDistance(include_background=True)

        def _val_func(engine, batch):
            pass

        engine = Engine(_val_func)
        sur_metric.attach(engine, "surface_distance")

        y_pred, y = TEST_SAMPLE_1
        sur_metric.update([y_pred, y])
        self.assertAlmostEqual(sur_metric.compute(), 4.17133, places=4)
        y_pred, y = TEST_SAMPLE_2
        sur_metric.update([y_pred, y])
        self.assertAlmostEqual(sur_metric.compute(), 2.08566, places=4)
        y_pred, y = TEST_SAMPLE_3
        sur_metric.update([y_pred, y])
        self.assertAlmostEqual(sur_metric.compute(), float("inf"))
        y_pred, y = TEST_SAMPLE_4
        sur_metric.update([y_pred, y])
        self.assertAlmostEqual(sur_metric.compute(), float("inf"))

    def test_shape_mismatch(self):
        sur_metric = SurfaceDistance(include_background=True)
        with self.assertRaises((AssertionError, ValueError)):
            y_pred = TEST_SAMPLE_1[0]
            y = torch.ones((1, 1, 10, 10, 10))
            sur_metric.update([y_pred, y])


if __name__ == "__main__":
    unittest.main()

[((52, 23, 52, 55), 'torch.zeros_like', 'torch.zeros_like', ({(52, 40, 52, 54): 'sampler_sphere'}, {}), '(sampler_sphere)', False, 'import torch\n'), ((38, 12, 38, 46), 'numpy.zeros', 'np.zeros', (), '', True, 'import numpy as np\n'), ((94, 4, 94, 19), 'unittest.main', 'unittest.main', ({}, {}), '()', False, 'import unittest\n'), ((64, 21, 64, 61), 'monai.handlers.SurfaceDistance', 'SurfaceDistance', (), '', False, 'from monai.handlers import SurfaceDistance\n'), ((69, 17, 69, 34), 'ignite.engine.Engine', 'Engine', ({(69, 24, 69, 33): '_val_func'}, {}), '(_val_func)', False, 'from ignite.engine import Engine\n'), ((86, 21, 86, 61), 'monai.handlers.SurfaceDistance', 'SurfaceDistance', (), '', False, 'from monai.handlers import SurfaceDistance\n'), ((89, 16, 89, 46), 'torch.ones', 'torch.ones', ({(89, 27, 89, 45): '(1, 1, 10, 10, 10)'}, {}), '((1, 1, 10, 10, 10))', False, 'import torch\n')]