import polars as pl
import pytest

from pipeline.utils.poi_counts import count_pois_per_postcode

POI_GROUPS = {
    "restaurants": ["Restaurant", "Fast Food"],
    "groceries": ["Supermarket"],
    "parks": ["Park"],
    "public_transport": ["Station"],
}


@pytest.fixture
def pois():
    """POIs clustered around two locations: central London and 10km away."""
    return pl.DataFrame(
        {
            "lat": [51.5074, 51.5075, 51.5080, 51.5076, 51.5073, 51.60],
            "lng": [-0.1278, -0.1280, -0.1275, -0.1279, -0.1277, -0.20],
            "category": [
                "Restaurant",
                "Fast Food",
                "Supermarket",
                "Park",
                "Station",
                "Restaurant",  # too far from any property
            ],
        }
    )


@pytest.fixture
def postcodes():
    """Two postcodes: one near central London, one far away."""
    return pl.DataFrame(
        {
            "postcode": ["EC1A 1BB", "ZZ99 9ZZ"],
            "lat": [51.5074, 55.0],
            "lon": [-0.1278, -3.0],
        }
    )


def test_counts_pois_within_radius(postcodes, pois):
    result = count_pois_per_postcode(postcodes, pois, groups=POI_GROUPS, radius_km=2.0)

    expected_cols = {f"{g}_2km" for g in POI_GROUPS}
    assert expected_cols.issubset(set(result.columns))

    # Result must be aligned to postcodes (2 rows)
    assert len(result) == 2

    ec1a = result.filter(pl.col("postcode") == "EC1A 1BB")
    assert ec1a["restaurants_2km"][0] == 2  # Restaurant + Fast Food
    assert ec1a["groceries_2km"][0] == 1  # Supermarket
    assert ec1a["parks_2km"][0] == 1  # Park
    assert ec1a["public_transport_2km"][0] == 1  # Station

    # Far-away postcode should have zero counts
    zz99 = result.filter(pl.col("postcode") == "ZZ99 9ZZ")
    for group in POI_GROUPS:
        assert zz99[f"{group}_2km"][0] == 0


def test_no_pois_returns_zeros(postcodes):
    empty_pois = pl.DataFrame(
        {
            "lat": pl.Series([], dtype=pl.Float64),
            "lng": pl.Series([], dtype=pl.Float64),
            "category": pl.Series([], dtype=pl.String),
        }
    )
    result = count_pois_per_postcode(postcodes, empty_pois, groups=POI_GROUPS, radius_km=2.0)

    for group in POI_GROUPS:
        col = f"{group}_2km"
        assert col in result.columns
        assert result[col].to_list() == [0, 0]


def test_custom_radius(pois):
    """A tiny radius should exclude POIs that are even slightly away."""
    postcodes = pl.DataFrame(
        {
            "postcode": ["EC1A 1BB"],
            "lat": [51.5074],
            "lon": [-0.1278],
        }
    )

    # 0.01 km = 10m — only the POI at the exact same location should match
    result = count_pois_per_postcode(postcodes, pois, groups=POI_GROUPS, radius_km=0.01)
    # The Restaurant at (51.5074, -0.1278) is at distance 0
    assert result["restaurants_0km"][0] >= 1
    # POIs >100m away should not be counted
    total = sum(result[f"{g}_0km"][0] for g in POI_GROUPS)
    assert total <= 2  # at most the co-located POIs