import numpy as np
import pytest
from numpy.testing import assert_allclose

from astropy.stats.spatial import RipleysKEstimator
from astropy.utils.misc import NumpyRNGContext

a = np.array([[1, 4], [2, 5], [3, 6]])
b = np.array([[-1, 1], [-2, 2], [-3, 3]])


@pytest.mark.parametrize("points, x_min, x_max", [(a, 0, 10), (b, -5, 5)])
def test_ripley_K_implementation(points, x_min, x_max):
    """
    Test against Ripley's K function implemented in R package `spatstat`
        +-+---------+---------+----------+---------+-+
      6 +                                          * +
        |                                            |
        |                                            |
    5.5 +                                            +
        |                                            |
        |                                            |
      5 +                     *                      +
        |                                            |
    4.5 +                                            +
        |                                            |
        |                                            |
      4 + *                                          +
        +-+---------+---------+----------+---------+-+
          1        1.5        2         2.5        3

        +-+---------+---------+----------+---------+-+
      3 + *                                          +
        |                                            |
        |                                            |
    2.5 +                                            +
        |                                            |
        |                                            |
      2 +                     *                      +
        |                                            |
    1.5 +                                            +
        |                                            |
        |                                            |
      1 +                                          * +
        +-+---------+---------+----------+---------+-+
         -3       -2.5       -2        -1.5       -1
    """

    area = 100
    r = np.linspace(0, 2.5, 5)
    Kest = RipleysKEstimator(
        area=area, x_min=x_min, y_min=x_min, x_max=x_max, y_max=x_max
    )

    ANS_NONE = np.array([0, 0, 0, 66.667, 66.667])
    assert_allclose(ANS_NONE, Kest(data=points, radii=r, mode="none"), atol=1e-3)

    ANS_TRANS = np.array([0, 0, 0, 82.304, 82.304])
    assert_allclose(
        ANS_TRANS, Kest(data=points, radii=r, mode="translation"), atol=1e-3
    )


with NumpyRNGContext(123):
    a = np.random.uniform(low=5, high=10, size=(100, 2))
    b = np.random.uniform(low=-5, high=-10, size=(100, 2))


@pytest.mark.parametrize("points", [a, b])
def test_ripley_uniform_property(points):
    # Ripley's K function without edge-correction converges to the area when
    # the number of points and the argument radii are large enough, i.e.,
    # K(x) --> area as x --> inf
    area = 50
    Kest = RipleysKEstimator(area=area)
    r = np.linspace(0, 20, 5)
    assert_allclose(area, Kest(data=points, radii=r, mode="none")[4])


with NumpyRNGContext(123):
    a = np.random.uniform(low=0, high=1, size=(500, 2))
    b = np.random.uniform(low=-1, high=0, size=(500, 2))


@pytest.mark.parametrize("points, low, high", [(a, 0, 1), (b, -1, 0)])
def test_ripley_large_density(points, low, high):
    Kest = RipleysKEstimator(area=1, x_min=low, x_max=high, y_min=low, y_max=high)
    r = np.linspace(0, 0.25, 25)
    Kpos = Kest.poisson(r)
    modes = ["ohser", "translation", "ripley"]
    for m in modes:
        Kest_r = Kest(data=points, radii=r, mode=m)
        assert_allclose(Kpos, Kest_r, atol=1e-1)


with NumpyRNGContext(123):
    a = np.random.uniform(low=5, high=10, size=(500, 2))
    b = np.random.uniform(low=-10, high=-5, size=(500, 2))


@pytest.mark.parametrize("points, low, high", [(a, 5, 10), (b, -10, -5)])
def test_ripley_modes(points, low, high):
    Kest = RipleysKEstimator(area=25, x_max=high, y_max=high, x_min=low, y_min=low)
    r = np.linspace(0, 1.2, 25)
    Kpos_mean = np.mean(Kest.poisson(r))
    modes = ["ohser", "translation", "ripley"]
    for m in modes:
        Kest_mean = np.mean(Kest(data=points, radii=r, mode=m))
        assert_allclose(Kpos_mean, Kest_mean, atol=1e-1, rtol=1e-1)


with NumpyRNGContext(123):
    a = np.random.uniform(low=0, high=1, size=(50, 2))
    b = np.random.uniform(low=-1, high=0, size=(50, 2))


@pytest.mark.parametrize("points, low, high", [(a, 0, 1), (b, -1, 0)])
def test_ripley_large_density_var_width(points, low, high):
    Kest = RipleysKEstimator(area=1, x_min=low, x_max=high, y_min=low, y_max=high)
    r = np.linspace(0, 0.25, 25)
    Kpos = Kest.poisson(r)
    Kest_r = Kest(data=points, radii=r, mode="var-width")
    assert_allclose(Kpos, Kest_r, atol=1e-1)


with NumpyRNGContext(123):
    a = np.random.uniform(low=5, high=10, size=(50, 2))
    b = np.random.uniform(low=-10, high=-5, size=(50, 2))


@pytest.mark.parametrize("points, low, high", [(a, 5, 10), (b, -10, -5)])
def test_ripley_var_width(points, low, high):
    Kest = RipleysKEstimator(area=25, x_max=high, y_max=high, x_min=low, y_min=low)
    r = np.linspace(0, 1.2, 25)
    Kest_ohser = np.mean(Kest(data=points, radii=r, mode="ohser"))
    Kest_var_width = np.mean(Kest(data=points, radii=r, mode="var-width"))
    assert_allclose(Kest_ohser, Kest_var_width, atol=1e-1, rtol=1e-1)