import numpy as np
import pytest
from typing import cast

from pyspat.core import MaskWindow, PPP, RectangleWindow, PolyWindow

def square_window(n=64, spacing=(1.0, 1.0), origin=(0.0, 0.0)):
    return MaskWindow(np.ones((n, n), bool), origin=origin, spacing=spacing)

def test_ppp_construct_and_validate():
    W = square_window(8)
    xy = np.array([(1.0, 1.0), (2.5, 3.5), (6.9, 0.1)])
    p = PPP(xy, W, unit="µm")
    assert len(p) == 3
    assert p.intensity() == pytest.approx(3 / W.area())
    with pytest.raises(ValueError):
        PPP(np.array([(100.0, 0.0)]), W)

def test_ppp_nndist_matches_bruteforce_small():
    W = square_window(64, spacing=(0.1, 0.1))
    rng = np.random.default_rng(123)
    xy = rng.random((50, 2)) * 6.0
    p = PPP(xy, W)
    d1 = p.nndist(k=1)
    diffs = xy[:, None, :] - xy[None, :, :]
    D = np.hypot(diffs[..., 0], diffs[..., 1])
    np.fill_diagonal(D, np.inf)
    d1_bf = D.min(axis=1)
    assert np.allclose(d1, d1_bf, rtol=0, atol=1e-12)

def test_ppp_thin_and_rescale():
    W = square_window(32)
    rng = np.random.default_rng(7)
    xy = rng.random((200, 2)) * 30.0
    p = PPP(xy, W, unit="µm")
    q = p.thin(0.4, rng=np.random.default_rng(123))
    assert 0 <= len(q) <= len(p)
    r = p.rescale(2.0, to_unit="2µm")
    mwin = cast(MaskWindow, r.window)
    assert np.allclose(r.xy, p.xy * 2.0)
    assert mwin.spacing == (2.0, 2.0)
    assert mwin.origin == (0.0, 0.0)
    assert r.intensity() == pytest.approx(len(p) / (W.area() * 4.0))

def test_ppp_rejects_outside_and_nonfinite():
    W = square_window(8)
    with pytest.raises(ValueError):
        PPP(np.array([[9.0, 1.0]]), W)
    with pytest.raises(ValueError):
        PPP(np.array([[np.nan, 0.0]]), W)
    with pytest.raises(ValueError):
        PPP(np.array([[0.0, 0.0, 0.0]]), W)

def test_ppp_marks_length_and_storage():
    W = square_window(8)
    xy = np.array([[1.0, 1.0], [2.0, 3.0]])
    marks = np.array([10, 20])
    p = PPP(xy, W, marks=marks, unit="µm")
    assert p.marks is not None and p.marks.shape == (2,)
    r = p.rescale(2.0)
    marks[0] = 999
    assert r.marks[0] == 10

def test_ppp_len_bbox_intensity_basic():
    W = square_window(10, spacing=(0.5, 0.5))
    xy = np.array([[0.1, 0.1], [2.0, 3.0], [4.9, 0.2]])
    p = PPP(xy, W)
    assert len(p) == 3
    xmin, ymin, xmax, ymax = p.bbox()
    assert (xmin, ymin) == (0.1, 0.1)
    assert (xmax, ymax) == (4.9, 3.0)
    assert p.intensity() == pytest.approx(3 / W.area())

def test_intensity_zero_area_raises():
    Wempty = MaskWindow(np.zeros((4, 4), bool))
    p = PPP(np.empty((0, 2)), Wempty)
    with pytest.raises(ValueError):
        p.intensity()

def test_thin_edges_and_rng_reproducible():
    W = square_window(32)
    xy = np.vstack(
        [np.linspace(0.1, 30.9, 200), np.linspace(0.2, 30.8, 200)]
    ).T
    p = PPP(xy, W)
    q1 = p.thin(0.3, rng=np.random.default_rng(123))
    q2 = p.thin(0.3, rng=np.random.default_rng(123))
    q3 = p.thin(0.3, rng=np.random.default_rng(124))
    assert np.array_equal(q1.xy, q2.xy)
    assert not np.array_equal(q1.xy, q3.xy)
    assert len(p.thin(0.0, rng=np.random.default_rng(0)).xy) == 0
    assert len(p.thin(1.0, rng=np.random.default_rng(0)).xy) == len(p)

def test_rescale_scales_geometry_and_contains():
    W = square_window(16)
    rng = np.random.default_rng(7)
    xy = rng.random((50, 2)) * 15.0
    p = PPP(xy, W, unit="µm")
    f = 2.0
    r = p.rescale(f, to_unit="2µm")
    assert np.allclose(r.xy, p.xy * f)
    assert r.window.contains(p.xy * f).all()
    assert r.intensity() == pytest.approx(p.intensity() / (f * f))

def test_nndist_shapes_and_ordering():
    W = square_window(64, spacing=(0.1, 0.1))
    rng = np.random.default_rng(1)
    xy = rng.random((40, 2)) * 5.0
    p = PPP(xy, W)
    d1 = p.nndist(1)
    d3 = p.nndist(3)
    assert d1.shape == (40,)
    assert d3.shape == (40, 3)
    assert np.all(np.diff(d3, axis=1) >= -1e-12)
    assert np.allclose(d1, d3[:, 0])

def test_nndist_vs_bruteforce_small():
    W = square_window(64, spacing=(0.5, 0.5))
    rng = np.random.default_rng(9)
    xy = rng.random((25, 2)) * 30.0
    p = PPP(xy, W)
    d2 = p.nndist(2)
    diffs = xy[:, None, :] - xy[None, :, :]
    D = np.hypot(diffs[..., 0], diffs[..., 1])
    np.fill_diagonal(D, np.inf)
    D_sorted = np.sort(D, axis=1)
    assert np.allclose(d2[:, 0], D_sorted[:, 0])
    assert np.allclose(d2[:, 1], D_sorted[:, 1])

def test_nndist_singleton_and_large_k():
    W = square_window(8)
    p_single = PPP(np.array([[1.0, 1.0]]), W)
    assert np.isinf(p_single.nndist(1)[0])
    xy = np.array(
        [[1.0, 1.0], [1.0, 2.0], [3.0, 1.0], [2.0, 3.0], [2.5, 2.1]]
    )
    p = PPP(xy, W)
    d = p.nndist(10)
    assert d.shape == (5, 10)
    assert np.all(np.isinf(d[:, 4:]))

def test_ppp_empty_pattern_behaviour():
    W = square_window(8)
    p = PPP(np.empty((0, 2)), W)
    assert len(p) == 0
    assert p.nndist(1).shape == (0,)
    assert p.nndist(3).shape == (0, 3)
    q = p.thin(0.5, rng=np.random.default_rng(0))
    r = p.rescale(2.0)
    assert len(q) == 0 and len(r) == 0
    with pytest.raises(ValueError):
        _ = p.bbox()

def test_ppp_duplicates_give_zero_nndist_and_sorted_cols():
    W = square_window(8)
    xy = np.array(
        [[1.0, 1.0], [1.0, 1.0], [2.0, 1.0], [4.0, 4.0]]
    )
    p = PPP(xy, W)
    d = p.nndist(3)
    assert np.all((d[:2, 0] == 0.0))
    assert np.all(np.diff(d, axis=1) >= -1e-12)

def test_thin_preserves_mark_alignment_and_length():
    W = square_window(32)
    rng = np.random.default_rng(123)
    xy = rng.random((200, 2)) * 30.0
    marks = np.arange(200)
    p = PPP(xy, W, marks=marks)
    seed = 999
    q = p.thin(0.4, rng=np.random.default_rng(seed))
    keep_mask = np.random.default_rng(seed).random(200) < 0.4
    assert q.marks is not None
    assert np.array_equal(q.marks, marks[keep_mask])
    assert len(q) == keep_mask.sum()

def test_rescale_with_nonzero_origin_and_anisotropic_spacing():
    origin = (10.0, -3.0)
    spacing = (0.5, 2.0)
    W = square_window(32, spacing=spacing, origin=origin)
    rng = np.random.default_rng(7)
    xy = np.column_stack(
        [
            origin[0] + rng.random(200) * spacing[0] * 30,
            origin[1] + rng.random(200) * spacing[1] * 30,
        ]
    )
    p = PPP(xy, W)
    f = 3.0
    r = p.rescale(f, to_unit="scaled")
    assert np.allclose(r.xy, p.xy * f)
    assert r.window.contains(p.xy * f).all()
    assert r.intensity() == pytest.approx(p.intensity() / (f * f))

def test_thin_yields_approximately_p_fraction():
    W = square_window(256)
    rng = np.random.default_rng(0)
    n = 10000
    xy = rng.random((n, 2)) * 200.0
    p = PPP(xy, W)
    p_keep = 0.25
    q = p.thin(p_keep, rng=np.random.default_rng(123))
    mean = n * p_keep
    sd = (n * p_keep * (1 - p_keep)) ** 0.5
    assert abs(len(q) - mean) <= 5 * sd

def test_nndist_pads_inf_when_k_exceeds_n_minus_1():
    W = square_window(16)
    xy = np.array([[0.5, 0.5], [1.5, 0.5], [0.5, 1.5]])
    p = PPP(xy, W)
    d = p.nndist(10)
    assert d.shape == (3, 10)
    assert np.all(np.isfinite(d[:, :2]))
    assert np.all(np.isinf(d[:, 2:]))

def rect_10x5():
    return RectangleWindow(0.0, 0.0, 10.0, 5.0)

def poly_square_10():
    pytest.importorskip("shapely")
    import shapely.geometry as sgeom
    return PolyWindow(sgeom.Polygon([(0, 0), (10, 0), (10, 10), (0, 10)]))

def poly_square_with_hole():
    pytest.importorskip("shapely")
    import shapely.geometry as sgeom
    outer = sgeom.Polygon([(0, 0), (10, 0), (10, 10), (0, 10)])
    hole = sgeom.Polygon([(4, 4), (6, 4), (6, 6), (4, 6)])
    return PolyWindow(sgeom.Polygon(outer.exterior.coords, [hole.exterior.coords]))

def sample_inside(win, n, seed=0):
    rng = np.random.default_rng(seed)
    return win.sample_uniform(n, rng)

def test_ppp_rectangle_construct_and_reject_outside():
    W = rect_10x5()
    xy = np.array([[1.0, 1.0], [9.9, 4.9], [5.0, 0.1]])
    p = PPP(xy, W)
    assert len(p) == 3
    with pytest.raises(ValueError):
        PPP(np.array([[10.0, 2.5]]), W)

def test_ppp_rectangle_intensity_bbox_and_rescale_invariants():
    W = rect_10x5()
    xy = sample_inside(W, 200, seed=1)
    p = PPP(xy, W, unit="µm")
    xmin, ymin, xmax, ymax = p.bbox()
    assert xmin >= 0.0 and ymin >= 0.0 and xmax <= 10.0 and ymax <= 5.0
    lam = p.intensity()
    f = 3.0
    r = p.rescale(f, to_unit="3µm")
    assert np.allclose(r.xy, p.xy * f)
    assert r.window.contains(r.xy).all()
    assert r.intensity() == pytest.approx(lam / (f * f))

def test_ppp_rectangle_nndist_matches_bruteforce():
    W = rect_10x5()
    xy = sample_inside(W, 60, seed=2)
    p = PPP(xy, W)
    d1 = p.nndist(1)
    diffs = xy[:, None, :] - xy[None, :, :]
    D = np.hypot(diffs[..., 0], diffs[..., 1])
    np.fill_diagonal(D, np.inf)
    d1_bf = D.min(axis=1)
    assert np.allclose(d1, d1_bf, atol=1e-12, rtol=0)

def test_ppp_rectangle_thin_reproducible_and_limits():
    W = rect_10x5()
    xy = sample_inside(W, 500, seed=3)
    p = PPP(xy, W)
    seed = 123
    q1 = p.thin(0.35, rng=np.random.default_rng(seed))
    q2 = p.thin(0.35, rng=np.random.default_rng(seed))
    q3 = p.thin(0.35, rng=np.random.default_rng(seed + 1))
    assert np.array_equal(q1.xy, q2.xy)
    assert not np.array_equal(q1.xy, q3.xy)
    assert len(p.thin(0.0, rng=np.random.default_rng(0))) == 0
    assert len(p.thin(1.0, rng=np.random.default_rng(0))) == len(p)

def test_ppp_poly_construct_and_reject_outside_and_hole():
    W = poly_square_10()
    xy = np.array([[1.0, 1.0], [9.9, 5.0], [5.0, 9.9]])
    p = PPP(xy, W)
    assert len(p) == 3
    with pytest.raises(ValueError):
        PPP(np.array([[11.0, 1.0]]), W)

    Wh = poly_square_with_hole()
    with pytest.raises(ValueError):
        PPP(np.array([[5.0, 5.0]]), Wh)

def test_ppp_poly_intensity_bbox_and_rescale_invariants():
    W = poly_square_10()
    xy = sample_inside(W, 300, seed=4)
    p = PPP(xy, W, unit="px")
    xmin, ymin, xmax, ymax = p.bbox()
    assert xmin >= 0.0 and ymin >= 0.0 and xmax <= 10.0 and ymax <= 10.0
    lam = p.intensity()
    f = 2.5
    r = p.rescale(f, to_unit="scaled")
    assert np.allclose(r.xy, p.xy * f)
    assert r.window.contains(r.xy).all()
    assert r.intensity() == pytest.approx(lam / (f * f))

def test_ppp_poly_nndist_matches_bruteforce_and_shapes():
    W = poly_square_10()
    xy = sample_inside(W, 45, seed=5)
    p = PPP(xy, W)
    d3 = p.nndist(3)
    assert d3.shape == (45, 3)
    assert np.all(np.diff(d3, axis=1) >= -1e-12)
    diffs = xy[:, None, :] - xy[None, :, :]
    D = np.hypot(diffs[..., 0], diffs[..., 1])
    np.fill_diagonal(D, np.inf)
    d1_bf = D.min(axis=1)
    assert np.allclose(d3[:, 0], d1_bf, atol=1e-12, rtol=0)

def test_ppp_poly_thin_reproducible_and_binomial_tolerance():
    W = poly_square_10()
    n = 5000
    xy = sample_inside(W, n, seed=6)
    p = PPP(xy, W)
    p_keep = 0.2
    rng = np.random.default_rng(777)
    q = p.thin(p_keep, rng=rng)
    mean = n * p_keep
    sd = (n * p_keep * (1 - p_keep)) ** 0.5
    assert abs(len(q) - mean) <= 5 * sd

def test_ppp_empty_with_poly_and_mask_have_same_api():
    pytest.importorskip("shapely")
    import shapely.geometry as sgeom

    Pempty = PolyWindow(sgeom.Polygon())
    p0 = PPP(np.empty((0, 2)), Pempty)
    assert len(p0) == 0
    assert p0.nndist(1).shape == (0,)
    with pytest.raises(ValueError):
        p0.intensity()

    Mempty = MaskWindow(np.zeros((3, 3), bool))
    pm = PPP(np.empty((0, 2)), Mempty)
    assert len(pm) == 0
    with pytest.raises(ValueError):
        pm.intensity()
