# SPDX-FileCopyrightText: Copyright © 2023 Idiap Research Institute <contact@idiap.ch>
#
# SPDX-License-Identifier: GPL-3.0-or-later
"""Tests for transforms."""

import numpy
import PIL.Image
import torchvision.transforms.functional as F # noqa: N812

from medbase.data.augmentations import ElasticDeformation


def test_elastic_deformation(datadir):
    # Get a raw sample without deformation
    data_file = str(datadir / "raw_without_elastic_deformation.png")
    raw_without_deformation = F.to_tensor(PIL.Image.open(data_file))

    # Elastic deforms the raw
    numpy.random.seed(seed=100)
    ed = ElasticDeformation()
    raw_deformed = ed(raw_without_deformation)

    # Get the same sample already deformed (with seed=100)
    data_file_2 = str(datadir / "raw_with_elastic_deformation.png")
    raw_2 = PIL.Image.open(data_file_2)

    # Compare both
    raw_deformed = (255 * numpy.asarray(raw_deformed)).astype(numpy.uint8)[
        0,
        :,
        :,
    ]
    raw_2 = numpy.asarray(raw_2)

    numpy.testing.assert_array_equal(raw_deformed, raw_2)