# Copyright (c) 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.

from __future__ import annotations

import unittest

import numpy as np
import torch

from monai.transforms import CastToType, Pad
from monai.utils import NumpyPadMode, PytorchPadMode
from tests.test_utils import dict_product


class TestPadMode(unittest.TestCase):
    def test_pad(self):
        expected_shapes = {3: (1, 15, 10), 4: (1, 10, 6, 7)}
        devices = ("cuda:0", "cpu") if torch.cuda.is_available() else ("cpu",)
        shapes = ((1, 10, 10), (1, 5, 6, 7))
        types = (float, int, np.uint8, np.int16, np.float32, bool)
        modes = list(PytorchPadMode) + list(NumpyPadMode)

        for params in dict_product(t=types, d=devices, s=shapes, m=modes):
            t = params["t"]
            d = params["d"]
            s = params["s"]
            m = params["m"]
            a = torch.rand(s)
            to_pad = [(0, 0), (2, 3)] if len(s) == 3 else [(0, 0), (2, 3), (0, 0), (0, 0)]
            out = Pad(to_pad=to_pad, mode=m)(CastToType(dtype=t)(a).to(d))
            self.assertEqual(out.shape, expected_shapes[len(s)])


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