"""England boundary polygon and bounding box for geographic filtering.

Uses shapely prepared geometry for fast single-point checks (osmium handlers)
and vectorized shapely.contains for batch checks (Polars DataFrames).
"""

import json
from pathlib import Path

import numpy as np
import shapely
from shapely.geometry import shape
from shapely.prepared import PreparedGeometry, prep

# Bounding box for fast pre-filtering before the precise polygon check
ENGLAND_BBOX_WEST = -6.45
ENGLAND_BBOX_SOUTH = 49.85
ENGLAND_BBOX_EAST = 1.77
ENGLAND_BBOX_NORTH = 55.82


def load_england_polygon(geojson_path: Path) -> PreparedGeometry:
    """Load England boundary as a prepared shapely polygon for fast contains checks."""
    with open(geojson_path) as f:
        data = json.load(f)
    geometry = shape(data["features"][0]["geometry"])
    return prep(geometry)


def in_england_mask(
    geojson_path: Path, lats: np.ndarray, lngs: np.ndarray
) -> np.ndarray:
    """Vectorized check: which (lat, lng) points are within England.

    Returns a boolean numpy array.
    """
    with open(geojson_path) as f:
        data = json.load(f)
    polygon = shape(data["features"][0]["geometry"])
    pts = shapely.points(lngs, lats)
    return shapely.contains(polygon, pts)