import numpy as np
import pytest

from pyspat.core import RectangleWindow, PolyWindow, MaskWindow, PPP
from pyspat.stats import Kest, Lest, pcf, Gest, Fest, Jest
from pyspat.simulate import rpoispp, rThomas, rSSI, rMaternI

HAS_SHAPELY = True
try:
    import shapely.geometry as sgeom
except Exception:
    HAS_SHAPELY = False
    sgeom = None


@pytest.fixture(scope="module")
def rect_w():
    return RectangleWindow(0.0, 0.0, 256.0, 256.0, dtb_shape=(256, 256))


@pytest.fixture(scope="module")
def poly_w():
    if not HAS_SHAPELY:
        pytest.skip("shapely not available")
    return PolyWindow(
        sgeom.Polygon([(0, 0), (256, 0), (256, 256), (0, 256)]),
        dtb_shape=(256, 256),
    )


def _csr_ppp(lmbda, W, rng):
    return rpoispp(lmbda, W, rng)


@pytest.fixture(scope="module")
def rect_ppp(rect_w):
    rng = np.random.default_rng(202401)
    return _csr_ppp(0.05, rect_w, rng)


@pytest.fixture(scope="module")
def poly_ppp(poly_w):
    rng = np.random.default_rng(202402)
    return _csr_ppp(0.05, poly_w, rng)


@pytest.fixture(scope="module")
def r_grid():
    r = np.linspace(2.0, 40.0, 25)
    mid = (r >= 5.0) & (r <= 30.0)
    return r, mid


@pytest.fixture(scope="module")
def edges_grid():
    edges = np.linspace(2.0, 40.0, 60)
    centres = 0.5 * (edges[:-1] + edges[1:])
    mid = (centres >= 5.0) & (centres <= 30.0)
    return edges, centres, mid


def test_rect_K_L_translation_match_CSR(rect_ppp, r_grid):
    r, mid = r_grid
    K_hat = Kest(rect_ppp, r, correction="translation")
    K_theory = np.pi * r**2
    rel = np.abs(K_hat[mid] / K_theory[mid] - 1.0)
    assert np.median(rel) < 0.07
    assert rel.max() < 0.15
    L_hat = Lest(rect_ppp, r, correction="translation")
    assert np.sqrt(((L_hat[mid] - r[mid]) ** 2).mean()) < 1.0


def test_rect_pcf_translation_near_unity(rect_ppp, edges_grid):
    edges, centres, mid = edges_grid
    rmid, g = pcf(rect_ppp, edges, correction="translation")
    assert rmid.shape == centres.shape
    assert np.all(np.isfinite(g))
    assert np.median(np.abs(g[mid] - 1.0)) < 0.08
    assert np.max(np.abs(g[mid] - 1.0)) < 0.20


def test_rect_translation_vs_border_K_close_midrange(rect_ppp):
    r = np.linspace(4.0, 25.0, 15)
    Kt = Kest(rect_ppp, r, correction="translation")
    Kb = Kest(rect_ppp, r, correction="border")
    rel = np.abs(Kt - Kb) / (np.pi * r**2)
    assert np.median(rel) < 0.10
    assert rel.max() < 0.25


def test_rect_G_F_J_border_match_CSR(rect_ppp):
    lam = 0.05
    r = np.linspace(0.5, 8.0, 60)
    G = Gest(rect_ppp, r, method="border")
    F = Fest(
        rect_ppp,
        r,
        nsamp=1500,
        method="border",
        rng=np.random.default_rng(123),
    )
    J = Jest(
        rect_ppp,
        r,
        method_G="border",
        method_F="border",
        nsamp_F=1500,
        rng=np.random.default_rng(123),
    )
    theory = 1.0 - np.exp(-lam * np.pi * r**2)
    mask = np.isfinite(G) & np.isfinite(F)
    assert mask.any()
    assert np.median(np.abs(G[mask] - theory[mask])) < 0.08
    assert np.median(np.abs(F[mask] - theory[mask])) < 0.08
    stable = (theory >= 0.1) & (theory <= 0.9)
    jmask = mask & stable & np.isfinite(J)
    assert jmask.any()
    assert np.max(np.abs(J[jmask] - 1.0)) < 0.20


def test_rect_K_translation_scales_under_rescale(rect_w):
    rng = np.random.default_rng(14)
    xy = rect_w.sample_uniform(4000, rng)
    p = PPP(xy, rect_w)
    r = np.linspace(2.0, 25.0, 20)
    K0 = Kest(p, r, correction="translation")
    c = 2.0
    p2 = p.rescale(c)
    K2 = Kest(p2, c * r, correction="translation")
    assert np.allclose(K2, (c * c) * K0, rtol=2e-2, atol=1e-8)


def test_rect_pcf_integrates_to_K_differences(rect_ppp, edges_grid):
    edges, centres, mid = edges_grid
    rmid, g = pcf(rect_ppp, edges, correction="translation")
    K_hi = Kest(rect_ppp, edges[1:], correction="translation")
    K_lo = Kest(rect_ppp, edges[:-1], correction="translation")
    dK = K_hi - K_lo
    dr = np.diff(edges)
    approx = 2.0 * np.pi * rmid * g * dr
    rel = np.abs(approx[mid] - dK[mid]) / np.maximum(1e-9, np.abs(dK[mid]))
    assert np.median(rel) < 0.15
    assert np.max(rel) < 0.35


def test_rect_border_nan_when_r_huge(rect_w):
    rng = np.random.default_rng(16)
    xy = rect_w.sample_uniform(200, rng)
    p = PPP(xy, rect_w)
    r = np.array([1e9])
    assert np.isnan(Kest(p, r, correction="border")[0])
    assert np.isnan(
        Fest(
            p,
            r,
            method="border",
            nsamp=1000,
            rng=np.random.default_rng(1),
        )[0]
    )


@pytest.mark.skipif(not HAS_SHAPELY, reason="shapely not available")
def test_poly_K_L_translation_match_CSR(poly_ppp, r_grid):
    r, mid = r_grid
    K_hat = Kest(poly_ppp, r, correction="translation")
    K_theory = np.pi * r**2
    rel = np.abs(K_hat[mid] / K_theory[mid] - 1.0)
    assert np.median(rel) < 0.07
    assert rel.max() < 0.15
    L_hat = Lest(poly_ppp, r, correction="translation")
    assert np.sqrt(((L_hat[mid] - r[mid]) ** 2).mean()) < 1.0


@pytest.mark.skipif(not HAS_SHAPELY, reason="shapely not available")
def test_poly_pcf_translation_near_unity(poly_ppp, edges_grid):
    edges, centres, mid = edges_grid
    rmid, g = pcf(poly_ppp, edges, correction="translation")
    assert rmid.shape == centres.shape
    assert np.all(np.isfinite(g))
    assert np.median(np.abs(g[mid] - 1.0)) < 0.08
    assert np.max(np.abs(g[mid] - 1.0)) < 0.20


@pytest.mark.skipif(not HAS_SHAPELY, reason="shapely not available")
def test_poly_translation_vs_border_K_close_midrange(poly_ppp):
    r = np.linspace(4.0, 25.0, 15)
    Kt = Kest(poly_ppp, r, correction="translation")
    Kb = Kest(poly_ppp, r, correction="border")
    rel = np.abs(Kt - Kb) / (np.pi * r**2)
    assert np.median(rel) < 0.10
    assert rel.max() < 0.30


@pytest.mark.skipif(not HAS_SHAPELY, reason="shapely not available")
def test_poly_G_F_J_border_match_CSR(poly_ppp):
    lam = 0.05
    r = np.linspace(0.5, 8.0, 60)
    G = Gest(poly_ppp, r, method="border")
    F = Fest(
        poly_ppp,
        r,
        nsamp=2500,
        method="border",
        rng=np.random.default_rng(123),
    )
    J = Jest(
        poly_ppp,
        r,
        method_G="border",
        method_F="border",
        nsamp_F=2500,
        rng=np.random.default_rng(123),
    )
    theory = 1.0 - np.exp(-lam * np.pi * r**2)
    mask = np.isfinite(G) & np.isfinite(F)
    assert mask.any()
    assert np.median(np.abs(G[mask] - theory[mask])) < 0.08
    assert np.median(np.abs(F[mask] - theory[mask])) < 0.08
    stable = (theory >= 0.1) & (theory <= 0.9)
    jmask = mask & stable & np.isfinite(J)
    assert jmask.any()
    assert np.max(np.abs(J[jmask] - 1.0)) < 0.15


@pytest.mark.skipif(not HAS_SHAPELY, reason="shapely not available")
def test_poly_K_translation_scales_under_rescale(poly_w):
    rng = np.random.default_rng(24)
    xy = poly_w.sample_uniform(3500, rng)
    p = PPP(xy, poly_w)
    r = np.linspace(2.0, 25.0, 20)
    K0 = Kest(p, r, correction="translation")
    c = 2.0
    p2 = p.rescale(c)
    K2 = Kest(p2, c * r, correction="translation")
    assert np.allclose(K2, (c * c) * K0, rtol=2e-2, atol=1e-8)


@pytest.mark.skipif(not HAS_SHAPELY, reason="shapely not available")
def test_poly_pcf_integrates_to_K_differences(poly_ppp, edges_grid):
    edges, centres, mid = edges_grid
    rmid, g = pcf(poly_ppp, edges, correction="translation")
    K_hi = Kest(poly_ppp, edges[1:], correction="translation")
    K_lo = Kest(poly_ppp, edges[:-1], correction="translation")
    dK = K_hi - K_lo
    dr = np.diff(edges)
    approx = 2.0 * np.pi * rmid * g * dr
    rel = np.abs(approx[mid] - dK[mid]) / np.maximum(1e-9, np.abs(dK[mid]))
    assert np.median(rel) < 0.15
    assert np.max(rel) < 0.35


@pytest.mark.skipif(not HAS_SHAPELY, reason="shapely not available")
def test_poly_border_nan_when_r_huge(poly_w):
    rng = np.random.default_rng(26)
    xy = poly_w.sample_uniform(200, rng)
    p = PPP(xy, poly_w)
    r = np.array([1e9])
    assert np.isnan(Kest(p, r, correction="border")[0])
    assert np.isnan(
        Fest(
            p,
            r,
            method="border",
            nsamp=1000,
            rng=np.random.default_rng(1),
        )[0]
    )


def test_translation_K_almost_agrees_across_window_types_for_same_points():
    rng = np.random.default_rng(30)
    L = 200
    M = MaskWindow(np.ones((L, L), bool), spacing=(1.0, 1.0))
    R = RectangleWindow(0.0, 0.0, float(L), float(L))

    xy = M.sample_uniform(3000, rng)
    pM = PPP(xy, M)
    pR = PPP(xy, R)

    r = np.linspace(2.0, 25.0, 20)
    KM = Kest(pM, r, correction="translation")
    KR = Kest(pR, r, correction="translation")

    rel_MR = np.abs(KM - KR) / np.maximum(1e-9, np.pi * r**2)
    assert np.median(rel_MR) < 0.05
    assert rel_MR.max() < 0.12

    if HAS_SHAPELY:
        import shapely.geometry as sgeom
        P = PolyWindow(sgeom.Polygon([(0, 0), (L, 0), (L, L), (0, L)]))
        pP = PPP(xy, P)
        KP = Kest(pP, r, correction="translation")

        rel_MP = np.abs(KM - KP) / np.maximum(1e-9, np.pi * r**2)
        rel_RP = np.abs(KR - KP) / np.maximum(1e-9, np.pi * r**2)
        assert np.median(rel_MP) < 0.05 and rel_MP.max() < 0.12
        assert np.median(rel_RP) < 0.03 and rel_RP.max() < 0.10

def test_plot_smoke_rect(rect_ppp):
    # Single quick figure per stat.
    import numpy as np
    r = np.linspace(2.0, 30.0, 40)
    edges = np.linspace(2.0, 30.0, 41)

    # K, L
    Kest(rect_ppp, r, correction="translation", showplot=True)
    Lest(rect_ppp, r, correction="translation", showplot=True)

    # g
    rmid, g = pcf(rect_ppp, edges, correction="translation", showplot=True)
    assert rmid.shape == g.shape

    # G, F, J
    G = Gest(rect_ppp, r, method="border", showplot=True)
    F = Fest(rect_ppp, r, nsamp=1200, method="border",
             rng=np.random.default_rng(123), showplot=True)
    J = Jest(rect_ppp, r, method_G="border", method_F="border",
             nsamp_F=1200, rng=np.random.default_rng(123), showplot=True)

    assert np.all(np.isfinite(Kest(rect_ppp, r, showplot=False)) | np.isnan(G))
    assert G.shape == F.shape == J.shape == r.shape


def test_mask_translation_cache_reuse():
    import numpy as np
    from pyspat.core import MaskWindow, PPP
    import pyspat.stats as S

    L = 128
    M = MaskWindow(np.ones((L, L), bool), spacing=(1.0, 1.0))
    rng = np.random.default_rng(0)
    xy = M.sample_uniform(1000, rng)
    p = PPP(xy, M)

    before = len(S._MTW_CACHE)
    _ = S.Kest(p, np.linspace(2.0, 20.0, 10), showplot=False)  # builds cache
    after1 = len(S._MTW_CACHE)
    _ = S.Kest(p, np.linspace(2.0, 20.0, 10), showplot=False)  # reuses cache
    after2 = len(S._MTW_CACHE)

    assert after1 == before + 1
    assert after2 == after1


def test_distance_to_boundary_idempotent(rect_w):
    img1 = rect_w.distance_to_boundary()
    img2 = rect_w.distance_to_boundary()
    assert np.array_equal(img1.values, img2.values)


class TestClusteredPatternStats:
    """Verify statistics correctly identify clustering."""

    @pytest.fixture
    def clustered_ppp(self, rect_w):
        rng = np.random.default_rng(555)
        return rThomas(0.0003, 5.0, 30.0, rect_w, rng)

    def test_K_exceeds_csr_theory(self, clustered_ppp):
        r = np.linspace(2.0, 25.0, 20)
        K = Kest(clustered_ppp, r, correction="translation")
        K_csr = np.pi * r ** 2
        # Clustered patterns have K > CSR K at small/mid r
        excess = K - K_csr
        assert np.mean(excess[r < 15]) > 0

    def test_L_exceeds_r(self, clustered_ppp):
        r = np.linspace(2.0, 25.0, 20)
        L = Lest(clustered_ppp, r, correction="translation")
        # L(r) > r indicates clustering
        assert np.mean(L[r < 15] - r[r < 15]) > 0

    def test_pcf_exceeds_unity(self, clustered_ppp):
        edges = np.linspace(2.0, 20.0, 25)
        rmid, g = pcf(clustered_ppp, edges, correction="translation")
        # g(r) > 1 at small r for clustered
        small = rmid < 10
        assert np.nanmean(g[small]) > 1.0

    def test_G_rises_faster_than_csr(self, clustered_ppp):
        r = np.linspace(0.5, 10.0, 30)
        G = Gest(clustered_ppp, r, method="border")
        lam = len(clustered_ppp) / clustered_ppp.window.area()
        G_csr = 1 - np.exp(-lam * np.pi * r ** 2)
        # Clustered: more close neighbors, G rises faster
        finite = np.isfinite(G)
        if finite.sum() > 5:
            assert np.mean(G[finite] - G_csr[finite]) > 0


class TestRegularPatternStats:
    """Verify statistics correctly identify regularity/inhibition."""

    @pytest.fixture
    def regular_ppp(self, rect_w):
        rng = np.random.default_rng(666)
        return rSSI(400, 8.0, rect_w, rng)

    def test_K_below_csr_at_small_r(self, regular_ppp):
        r = np.linspace(2.0, 25.0, 20)
        K = Kest(regular_ppp, r, correction="translation")
        K_csr = np.pi * r ** 2
        # Hardcore: K = 0 for r < hardcore distance
        small = r < 8
        assert np.all(K[small] < K_csr[small])

    def test_L_below_r_at_small_r(self, regular_ppp):
        r = np.linspace(2.0, 25.0, 20)
        L = Lest(regular_ppp, r, correction="translation")
        small = r < 8
        assert np.all(L[small] < r[small])

    def test_pcf_zero_below_hardcore(self, regular_ppp):
        edges = np.linspace(1.0, 20.0, 30)
        rmid, g = pcf(regular_ppp, edges, correction="translation")
        # g(r) ~ 0 for r < hardcore
        below_hc = rmid < 7.5
        assert np.nanmean(g[below_hc]) < 0.1

    def test_G_zero_below_hardcore(self, regular_ppp):
        r = np.linspace(0.5, 15.0, 40)
        G = Gest(regular_ppp, r, method="border")
        # No neighbors within hardcore distance
        below_hc = r < 7.5
        finite = np.isfinite(G)
        assert np.all(G[below_hc & finite] < 0.05)


class TestMaternHardcoreStats:
    """Test statistics on Matern I thinned patterns."""

    @pytest.fixture
    def matern_ppp(self, rect_w):
        rng = np.random.default_rng(777)
        return rMaternI(0.02, 6.0, rect_w, rng)

    def test_nndist_respects_hardcore(self, matern_ppp):
        if len(matern_ppp) > 1:
            d1 = matern_ppp.nndist(1)
            assert np.all(d1 > 6.0 - 1e-9)

    def test_K_deficit_at_small_r(self, matern_ppp):
        r = np.linspace(2.0, 20.0, 15)
        K = Kest(matern_ppp, r, correction="translation")
        K_csr = np.pi * r ** 2
        small = r < 6
        assert np.all(K[small] < K_csr[small] * 0.5)