"""
Tests for point process simulation functions.
"""
import numpy as np
import pytest

from pyspat.core import RectangleWindow, MaskWindow, PPP, Image
from pyspat.simulate import (
    rpoispp, rpoispp_inhom, rThomas, rMatClust, rSSI, rMaternI, rMaternII
)

HAS_SHAPELY = True
try:
    import shapely.geometry as sgeom
    from pyspat.core import PolyWindow
except ImportError:
    HAS_SHAPELY = False
    sgeom = None
    PolyWindow = None



@pytest.fixture
def rect_window():
    return RectangleWindow(0.0, 0.0, 100.0, 100.0)


@pytest.fixture
def mask_window():
    mask = np.ones((100, 100), bool)
    return MaskWindow(mask, origin=(0.0, 0.0), spacing=(1.0, 1.0))


@pytest.fixture
def poly_window():
    if not HAS_SHAPELY:
        pytest.skip("shapely not available")
    return PolyWindow(sgeom.Polygon([(0, 0), (100, 0), (100, 100), (0, 100)]))


class TestRpoispp:

    def test_returns_ppp(self, rect_window):
        p = rpoispp(0.01, rect_window, rng=np.random.default_rng(0))
        assert isinstance(p, PPP)
        assert p.window is rect_window

    def test_mean_count_matches_intensity(self, rect_window):
        intensity = 0.05
        area = rect_window.area()
        expected_mean = intensity * area
        rng = np.random.default_rng(42)
        counts = [len(rpoispp(intensity, rect_window, rng)) for _ in range(200)]
        sample_mean = np.mean(counts)
        sample_std = np.std(counts)
        # Mean should be close to expected (within ~3 standard errors)
        assert abs(sample_mean - expected_mean) < 3 * sample_std / np.sqrt(200)

    def test_variance_matches_poisson(self, rect_window):
        intensity = 0.05
        area = rect_window.area()
        expected_var = intensity * area  # Poisson variance = mean
        rng = np.random.default_rng(123)
        counts = [len(rpoispp(intensity, rect_window, rng)) for _ in range(500)]
        sample_var = np.var(counts)
        # Variance should be close to expected
        assert abs(sample_var - expected_var) / expected_var < 0.2

    def test_points_inside_window(self, rect_window):
        rng = np.random.default_rng(0)
        for _ in range(10):
            p = rpoispp(0.1, rect_window, rng)
            assert rect_window.contains(p.xy).all()

    def test_reproducible_with_seed(self, rect_window):
        p1 = rpoispp(0.05, rect_window, rng=np.random.default_rng(999))
        p2 = rpoispp(0.05, rect_window, rng=np.random.default_rng(999))
        assert np.array_equal(p1.xy, p2.xy)

    def test_different_seeds_different_results(self, rect_window):
        p1 = rpoispp(0.05, rect_window, rng=np.random.default_rng(1))
        p2 = rpoispp(0.05, rect_window, rng=np.random.default_rng(2))
        assert not np.array_equal(p1.xy, p2.xy)

    def test_zero_intensity_raises(self, rect_window):
        with pytest.raises(ValueError):
            rpoispp(0.0, rect_window)

    def test_negative_intensity_raises(self, rect_window):
        with pytest.raises(ValueError):
            rpoispp(-0.1, rect_window)

    def test_very_low_intensity_mostly_empty(self, rect_window):
        rng = np.random.default_rng(0)
        counts = [len(rpoispp(1e-6, rect_window, rng)) for _ in range(100)]
        assert np.mean(counts) < 1  # Expected ~0.01 points

    def test_works_with_mask_window(self, mask_window):
        p = rpoispp(0.01, mask_window, rng=np.random.default_rng(0))
        assert isinstance(p, PPP)
        assert mask_window.contains(p.xy).all()

    @pytest.mark.skipif(not HAS_SHAPELY, reason="shapely not available")
    def test_works_with_poly_window(self, poly_window):
        p = rpoispp(0.01, poly_window, rng=np.random.default_rng(0))
        assert isinstance(p, PPP)
        assert poly_window.contains(p.xy).all()



class TestRpoisppInhom:

    def test_with_constant_function_matches_homogeneous(self, rect_window):
        lam = 0.05
        rng1 = np.random.default_rng(42)
        rng2 = np.random.default_rng(42)

        # Constant intensity function
        counts_inhom = [
            len(rpoispp_inhom(lambda xy: np.full(len(xy), lam), rect_window, lmax=lam, rng=rng1))
            for _ in range(100)
        ]
        counts_hom = [len(rpoispp(lam, rect_window, rng2)) for _ in range(100)]

        # Means should be similar
        assert abs(np.mean(counts_inhom) - np.mean(counts_hom)) < 50

    def test_with_intensity_image(self, rect_window):
        # Create intensity image: higher in center
        vals = np.zeros((100, 100))
        vals[30:70, 30:70] = 0.1
        img = Image(origin=(0.0, 0.0), spacing=(1.0, 1.0), values=vals)

        rng = np.random.default_rng(0)
        p = rpoispp_inhom(img, rect_window, rng=rng)

        # Most points should be in center region
        in_center = ((p.xy[:, 0] >= 30) & (p.xy[:, 0] < 70) &
                     (p.xy[:, 1] >= 30) & (p.xy[:, 1] < 70))
        assert in_center.mean() > 0.8

    def test_gradient_intensity(self, rect_window):
        # Intensity increases with x
        def intensity_func(xy):
            return 0.001 * xy[:, 0]  # 0 at x=0, 0.1 at x=100

        rng = np.random.default_rng(123)
        p = rpoispp_inhom(intensity_func, rect_window, lmax=0.1, rng=rng)

        # Mean x should be > 50 (skewed right)
        if len(p) > 10:
            assert np.mean(p.xy[:, 0]) > 50

    def test_lmax_required_for_callable(self, rect_window):
        with pytest.raises(ValueError, match="lmax must be provided"):
            rpoispp_inhom(lambda xy: np.ones(len(xy)), rect_window)

    def test_zero_lmax_raises(self, rect_window):
        vals = np.zeros((10, 10))
        img = Image(origin=(0.0, 0.0), spacing=(10.0, 10.0), values=vals)
        with pytest.raises(ValueError, match="lmax must be positive"):
            rpoispp_inhom(img, rect_window)

    def test_reproducible_with_seed(self, rect_window):
        vals = np.ones((10, 10)) * 0.05
        img = Image(origin=(0.0, 0.0), spacing=(10.0, 10.0), values=vals)

        p1 = rpoispp_inhom(img, rect_window, rng=np.random.default_rng(42))
        p2 = rpoispp_inhom(img, rect_window, rng=np.random.default_rng(42))
        assert np.array_equal(p1.xy, p2.xy)


class TestRThomas:

    def test_returns_ppp(self, rect_window):
        p = rThomas(0.001, 5.0, 10.0, rect_window, rng=np.random.default_rng(0))
        assert isinstance(p, PPP)

    def test_expected_intensity(self, rect_window):
        kappa, sigma, mu = 0.001, 5.0, 10.0
        expected_intensity = kappa * mu
        rng = np.random.default_rng(0)

        intensities = []
        for _ in range(100):
            p = rThomas(kappa, sigma, mu, rect_window, rng=rng)
            if rect_window.area() > 0:
                intensities.append(len(p) / rect_window.area())

        mean_intensity = np.mean(intensities)
        # Should be close to expected (some edge effects expected)
        assert abs(mean_intensity - expected_intensity) / expected_intensity < 0.3

    def test_points_are_clustered(self, rect_window):
        # High clustering: small sigma, high mu
        rng = np.random.default_rng(42)
        p = rThomas(0.0005, 2.0, 20.0, rect_window, rng=rng)

        if len(p) > 10:
            # Nearest-neighbor distances should be smaller than CSR
            d1 = p.nndist(1)
            mean_nnd = np.mean(d1)
            # For CSR, E[nnd] ≈ 0.5 / sqrt(lambda)
            lam = len(p) / rect_window.area()
            csr_nnd = 0.5 / np.sqrt(lam) if lam > 0 else np.inf
            assert mean_nnd < csr_nnd  # Clustered => smaller nnd

    def test_invalid_parameters_raise(self, rect_window):
        with pytest.raises(ValueError):
            rThomas(0.0, 5.0, 10.0, rect_window)  # kappa = 0
        with pytest.raises(ValueError):
            rThomas(0.001, 0.0, 10.0, rect_window)  # sigma = 0
        with pytest.raises(ValueError):
            rThomas(0.001, 5.0, 0.0, rect_window)  # mu = 0
        with pytest.raises(ValueError):
            rThomas(-0.001, 5.0, 10.0, rect_window)  # negative

    def test_reproducible(self, rect_window):
        p1 = rThomas(0.001, 5.0, 10.0, rect_window, rng=np.random.default_rng(99))
        p2 = rThomas(0.001, 5.0, 10.0, rect_window, rng=np.random.default_rng(99))
        assert np.array_equal(p1.xy, p2.xy)

    def test_points_inside_window(self, rect_window):
        rng = np.random.default_rng(0)
        for _ in range(10):
            p = rThomas(0.001, 5.0, 10.0, rect_window, rng=rng)
            assert rect_window.contains(p.xy).all()


class TestRMatClust:

    def test_returns_ppp(self, rect_window):
        p = rMatClust(0.001, 10.0, 10.0, rect_window, rng=np.random.default_rng(0))
        assert isinstance(p, PPP)

    def test_expected_intensity(self, rect_window):
        kappa, r, mu = 0.001, 10.0, 10.0
        expected_intensity = kappa * mu
        rng = np.random.default_rng(0)

        intensities = []
        for _ in range(100):
            p = rMatClust(kappa, r, mu, rect_window, rng=rng)
            if rect_window.area() > 0:
                intensities.append(len(p) / rect_window.area())

        mean_intensity = np.mean(intensities)
        assert abs(mean_intensity - expected_intensity) / expected_intensity < 0.3

    def test_offspring_within_radius(self, rect_window):
        # Large expand to avoid edge effects
        # Since offspring are uniform in disc, max dist from parent = r
        # We can't directly test this without tracking parents, but we can
        # check clustering behavior
        rng = np.random.default_rng(42)
        p = rMatClust(0.0002, 5.0, 50.0, rect_window, rng=rng)

        if len(p) > 10:
            d1 = p.nndist(1)
            # Many points should have very close neighbors (in same cluster)
            assert np.median(d1) < 5.0

    def test_invalid_parameters_raise(self, rect_window):
        with pytest.raises(ValueError):
            rMatClust(0.0, 10.0, 10.0, rect_window)
        with pytest.raises(ValueError):
            rMatClust(0.001, 0.0, 10.0, rect_window)
        with pytest.raises(ValueError):
            rMatClust(0.001, 10.0, 0.0, rect_window)

    def test_reproducible(self, rect_window):
        p1 = rMatClust(0.001, 10.0, 10.0, rect_window, rng=np.random.default_rng(77))
        p2 = rMatClust(0.001, 10.0, 10.0, rect_window, rng=np.random.default_rng(77))
        assert np.array_equal(p1.xy, p2.xy)


class TestRSSI:

    def test_returns_ppp(self, rect_window):
        p = rSSI(50, 3.0, rect_window, rng=np.random.default_rng(0))
        assert isinstance(p, PPP)

    def test_minimum_distance_respected(self, rect_window):
        r = 5.0
        rng = np.random.default_rng(0)
        p = rSSI(100, r, rect_window, rng=rng)

        if len(p) > 1:
            d1 = p.nndist(1)
            assert np.all(d1 > r - 1e-10)

    def test_exact_count_when_possible(self, rect_window):
        # Small n, small r => should achieve target
        rng = np.random.default_rng(42)
        p = rSSI(10, 1.0, rect_window, rng=rng)
        assert len(p) == 10

    def test_may_not_reach_target_when_packed(self, rect_window):
        # Very high density request with large hardcore radius
        # 100x100 window, r=20 => max ~25 circles fit
        rng = np.random.default_rng(0)
        p = rSSI(1000, 20.0, rect_window, rng=rng, max_attempts=10000)
        assert len(p) < 1000
        assert len(p) > 0

    def test_zero_radius_gives_n_points(self, rect_window):
        rng = np.random.default_rng(0)
        p = rSSI(50, 0.0, rect_window, rng=rng)
        assert len(p) == 50

    def test_zero_n_returns_empty(self, rect_window):
        p = rSSI(0, 5.0, rect_window)
        assert len(p) == 0

    def test_invalid_parameters_raise(self, rect_window):
        with pytest.raises(ValueError):
            rSSI(-1, 5.0, rect_window)
        with pytest.raises(ValueError):
            rSSI(10, -1.0, rect_window)

    def test_reproducible(self, rect_window):
        p1 = rSSI(30, 5.0, rect_window, rng=np.random.default_rng(123))
        p2 = rSSI(30, 5.0, rect_window, rng=np.random.default_rng(123))
        assert np.array_equal(p1.xy, p2.xy)


class TestRMaternI:

    def test_returns_ppp(self, rect_window):
        p = rMaternI(0.05, 5.0, rect_window, rng=np.random.default_rng(0))
        assert isinstance(p, PPP)

    def test_minimum_distance_respected(self, rect_window):
        r = 3.0
        rng = np.random.default_rng(42)
        p = rMaternI(0.1, r, rect_window, rng=rng)

        if len(p) > 1:
            d1 = p.nndist(1)
            assert np.all(d1 > r - 1e-10)

    def test_deletes_conflicting_pairs(self, rect_window):
        # High intensity, large r => many deletions
        rng = np.random.default_rng(0)
        p = rMaternI(0.5, 10.0, rect_window, rng=rng)

        # Should have far fewer points than 0.5 * 10000 = 5000 expected
        assert len(p) < 500

    def test_zero_radius_preserves_all(self, rect_window):
        rng1 = np.random.default_rng(0)
        rng2 = np.random.default_rng(0)
        p_all = rpoispp(0.05, rect_window, rng1)
        p_matern = rMaternI(0.05, 0.0, rect_window, rng2)
        assert len(p_matern) == len(p_all)

    def test_invalid_parameters_raise(self, rect_window):
        with pytest.raises(ValueError):
            rMaternI(0.0, 5.0, rect_window)
        with pytest.raises(ValueError):
            rMaternI(0.05, -1.0, rect_window)

    def test_reproducible(self, rect_window):
        p1 = rMaternI(0.05, 5.0, rect_window, rng=np.random.default_rng(55))
        p2 = rMaternI(0.05, 5.0, rect_window, rng=np.random.default_rng(55))
        assert np.array_equal(p1.xy, p2.xy)


class TestRMaternII:

    def test_returns_ppp(self, rect_window):
        p = rMaternII(0.05, 5.0, rect_window, rng=np.random.default_rng(0))
        assert isinstance(p, PPP)

    def test_minimum_distance_respected(self, rect_window):
        r = 3.0
        rng = np.random.default_rng(42)
        p = rMaternII(0.1, r, rect_window, rng=rng)

        if len(p) > 1:
            d1 = p.nndist(1)
            assert np.all(d1 > r - 1e-10)

    def test_retains_more_than_matern_I(self, rect_window):
        # Matérn II should typically keep more points than Matérn I
        rng1 = np.random.default_rng(0)
        rng2 = np.random.default_rng(0)

        counts_I = [len(rMaternI(0.1, 5.0, rect_window, rng1)) for _ in range(50)]
        counts_II = [len(rMaternII(0.1, 5.0, rect_window, rng2)) for _ in range(50)]

        # On average, II should retain more
        assert np.mean(counts_II) >= np.mean(counts_I) * 0.9

    def test_zero_radius_preserves_all(self, rect_window):
        rng1 = np.random.default_rng(0)
        rng2 = np.random.default_rng(0)
        p_all = rpoispp(0.05, rect_window, rng1)
        p_matern = rMaternII(0.05, 0.0, rect_window, rng2)
        assert len(p_matern) == len(p_all)

    def test_invalid_parameters_raise(self, rect_window):
        with pytest.raises(ValueError):
            rMaternII(0.0, 5.0, rect_window)
        with pytest.raises(ValueError):
            rMaternII(0.05, -1.0, rect_window)

    def test_reproducible(self, rect_window):
        p1 = rMaternII(0.05, 5.0, rect_window, rng=np.random.default_rng(88))
        p2 = rMaternII(0.05, 5.0, rect_window, rng=np.random.default_rng(88))
        assert np.array_equal(p1.xy, p2.xy)


class TestProcessComparisons:

    def test_csr_vs_clustered_K(self, rect_window):
        """Clustered processes should have K > CSR K at mid-range r."""
        from pyspat.stats import Kest

        rng = np.random.default_rng(42)
        r = np.linspace(2.0, 20.0, 15)

        # CSR
        p_csr = rpoispp(0.01, rect_window, rng)
        K_csr = Kest(p_csr, r, correction="translation")

        # Clustered (Thomas)
        p_clust = rThomas(0.0005, 3.0, 20.0, rect_window, rng)
        K_clust = Kest(p_clust, r, correction="translation")

        # Clustered should have higher K (more pairs at short distances)
        # At small r, difference should be clear
        mid_idx = len(r) // 2
        assert np.mean(K_clust[:mid_idx]) > np.mean(K_csr[:mid_idx])

    def test_csr_vs_regular_K(self, rect_window):
        """Regular (hardcore) processes should have K < CSR K at small r."""
        from pyspat.stats import Kest

        rng = np.random.default_rng(42)
        r = np.linspace(2.0, 20.0, 15)

        # CSR
        p_csr = rpoispp(0.01, rect_window, rng)
        K_csr = Kest(p_csr, r, correction="translation")

        # Regular (SSI)
        p_reg = rSSI(int(0.01 * rect_window.area()), 5.0, rect_window, rng)
        K_reg = Kest(p_reg, r, correction="translation")

        # At r < hardcore distance, K_reg should be near 0
        small_r_idx = r < 5.0
        assert np.mean(K_reg[small_r_idx]) < np.mean(K_csr[small_r_idx])


class TestEdgeCases:

    def test_tiny_window(self):
        tiny = RectangleWindow(0.0, 0.0, 1.0, 1.0)
        p = rpoispp(1.0, tiny, rng=np.random.default_rng(0))
        assert isinstance(p, PPP)

    def test_anisotropic_mask_window(self):
        mask = np.ones((50, 200), bool)
        w = MaskWindow(mask, origin=(0.0, 0.0), spacing=(0.5, 2.0))
        p = rpoispp(0.01, w, rng=np.random.default_rng(0))
        assert w.contains(p.xy).all()

    @pytest.mark.skipif(not HAS_SHAPELY, reason="shapely not available")
    def test_poly_with_hole(self):
        outer = sgeom.Polygon([(0, 0), (100, 0), (100, 100), (0, 100)])
        hole = sgeom.Polygon([(40, 40), (60, 40), (60, 60), (40, 60)])
        poly = sgeom.Polygon(outer.exterior.coords, [hole.exterior.coords])
        w = PolyWindow(poly)

        p = rpoispp(0.01, w, rng=np.random.default_rng(0))
        assert w.contains(p.xy).all()

        # No points should be in hole region
        in_hole = ((p.xy[:, 0] >= 40) & (p.xy[:, 0] <= 60) &
                   (p.xy[:, 1] >= 40) & (p.xy[:, 1] <= 60))
        assert not in_hole.any()

    def test_very_high_intensity(self, rect_window):
        # Should handle many points
        p = rpoispp(10.0, rect_window, rng=np.random.default_rng(0))
        assert len(p) > 50000  # Expected ~100000

    def test_thomas_with_large_sigma(self, rect_window):
        # Large sigma relative to window => nearly homogeneous
        rng = np.random.default_rng(0)
        p = rThomas(0.001, 100.0, 10.0, rect_window, rng=rng)
        assert isinstance(p, PPP)