import polars as pl

from pipeline.transform.poi_proximity import (
    _build_poi_category_groups,
    _dynamic_poi_metric_renames,
)


def test_dynamic_poi_groups_include_requested_categories_only() -> None:
    pois = pl.DataFrame(
        {
            "group": (
                ["Public Transport"] * 2
                + ["Leisure"] * 2
                + ["Groceries"] * 101
                + ["Groceries"] * 100
                + ["Leisure"] * 10
                + ["Education"] * 200
                + ["Health"] * 200
            ),
            "category": (
                ["Rail station", "Bus stop"]
                + ["Café", "Restaurant"]
                + ["Tesco"] * 101
                + ["Waitrose"] * 100
                + ["Park"] * 10
                + ["School"] * 200
                + ["Pharmacy"] * 200
            ),
            "lat": [51.5] * 615,
            "lng": [-0.1] * 615,
        }
    )

    groups, display_names = _build_poi_category_groups(pois)

    assert set(display_names.values()) == {
        "Bus stop",
        "Café",
        "Rail station",
        "Restaurant",
        "Tesco",
    }
    assert "poi_waitrose" not in groups
    assert "poi_park" not in groups
    assert "poi_school" not in groups
    assert "poi_pharmacy" not in groups


def test_dynamic_poi_metric_renames_support_park_count_options() -> None:
    assert _dynamic_poi_metric_renames({"parks": "Park"}) == {
        "parks_nearest_km": "Distance to nearest amenity (Park) (km)",
        "parks_2km": "Number of amenities (Park) within 2km",
        "parks_5km": "Number of amenities (Park) within 5km",
    }