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

import numpy
import PIL.Image

from ptbench.data.image_utils import (
    RemoveBlackBorders,
    SingleAutoLevel16to8,
    load_pil,
)


def test_remove_black_borders(datadir):
    # Get a raw sample with black border
    data_file = str(datadir / "raw_with_black_border.png")
    raw_with_black_border = PIL.Image.open(data_file)

    # Remove the black border
    rbb = RemoveBlackBorders()
    raw_rbb_removed = rbb(raw_with_black_border)

    # Get the same sample without black border
    data_file_2 = str(datadir / "raw_without_black_border.png")
    raw_without_black_border = PIL.Image.open(data_file_2)

    # Compare both
    raw_rbb_removed = numpy.asarray(raw_rbb_removed)
    raw_without_black_border = numpy.asarray(raw_without_black_border)

    numpy.testing.assert_array_equal(raw_without_black_border, raw_rbb_removed)


def test_load_pil_16bit(datadir):
    # If the ratio is higher 0.5, image is probably clipped
    Level16to8 = SingleAutoLevel16to8()

    data_file = str(datadir / "16bits.png")
    image = numpy.array(Level16to8(load_pil(data_file)))

    count_pixels = numpy.count_nonzero(image)
    count_max_value = numpy.count_nonzero(image == image.max())

    assert count_max_value / count_pixels < 0.5

    # It should not do anything to an image already in 8 bits
    data_file = str(datadir / "raw_without_black_border.png")
    img_loaded = load_pil(data_file)

    original_8bits = numpy.array(img_loaded)
    leveled_8bits = numpy.array(Level16to8(img_loaded))

    numpy.testing.assert_array_equal(original_8bits, leveled_8bits)