from __future__ import annotations

from collections.abc import Sequence

from hypothesis import given
from hypothesis.strategies import floats
from hypothesis.strategies import integers
from numpy import array
from numpy import isfinite
from numpy import isnan
from numpy import nan
from numpy.testing import assert_allclose
from numpy.testing import assert_equal
from pytest import mark
from pytest import param

from utilities.hypothesis import float_arrays
from utilities.numpy import NDArrayF1
from utilities.numpy import is_between_or_nan
from utilities.scipy import ppf


class TestPPF:
    @mark.parametrize(
        ("values", "expected"),
        [
            param([], []),
            param([0.0], [0.0]),
            param([0.0, 0.1], [-1.0, 1.0]),
            param([0.0, 0.1, 0.2], [-1.0, 0.0, 1.0]),
            param([0.0, 0.1, 0.2, 0.3], [-1.0, -0.2891889, 0.2891889, 1.0]),
            param([nan], [nan]),
            param([0.0, nan, 0.1], [-1.0, nan, 1.0]),
        ],
    )
    def test_examples(
        self, values: Sequence[float], expected: Sequence[float]
    ) -> None:
        result = ppf(array(values, dtype=float), 1.0)
        assert_allclose(result, array(expected, dtype=float))

    @given(
        array=float_arrays(
            shape=integers(0, 10),
            min_value=-10.0,
            max_value=10.0,
            allow_nan=True,
        ),
        cutoff=floats(0.0, 10.0),
    )
    def test_main(self, array: NDArrayF1, cutoff: float) -> None:
        result = ppf(array, cutoff)
        assert_equal(isfinite(result), isfinite(array))
        assert_equal(isnan(result), isnan(array))
        assert is_between_or_nan(result, -cutoff, cutoff).all()
