"""Snap an arbitrary lat/lon to a stable ~4-square-mile grid cell.

The grid is defined by fixed latitude rows (~2 miles tall). Within each row the
longitude step is scaled by cos(latitude) so cells stay roughly square (~4 sq mi)
at every latitude instead of getting skinny toward the poles. Cell ids are
deterministic, so the same physical location always maps to the same cache file.
Coverage is worldwide: any lat in [-90, 90] and any longitude (wrapped into
[-180, 180)) maps to a cell.
"""
import math

# 1 degree of latitude ~= 69 miles. ~2 miles -> ~0.029 deg gives a ~4 sq mi cell.
LAT_STEP = 1.0 / 34.5  # ~= 0.02899 deg  (~2.0 miles)


def _lon_step(center_lat: float) -> float:
    """Longitude degrees that span ~2 miles at the given latitude."""
    c = math.cos(math.radians(center_lat))
    c = max(c, 0.05)  # clamp near the poles to avoid a blow-up
    return LAT_STEP / c


def _cell(i: int, j: int) -> dict:
    """Build the cell dict for grid indices (i, j). Shared by snap()/from_id()
    so an id always rebuilds to the exact same cell."""
    center_lat = (i + 0.5) * LAT_STEP
    lon_step = _lon_step(center_lat)
    center_lon = (j + 0.5) * lon_step
    # Keep the reported center a valid coordinate for the upstream weather and
    # geocoding APIs: the topmost row's center overshoots the pole, and a row's
    # outermost cells can have centers just past the antimeridian on either side.
    center_lat = min(max(center_lat, -90.0), 90.0)
    if center_lon > 180.0:
        center_lon -= 360.0
    elif center_lon < -180.0:
        center_lon += 360.0

    # Approximate cell dimensions in miles for display.
    height_mi = LAT_STEP * 69.0
    width_mi = lon_step * 69.0 * math.cos(math.radians(center_lat))

    return {
        "id": f"{i}_{j}",
        "center_lat": round(center_lat, 5),
        "center_lon": round(center_lon, 5),
        "lat_step": LAT_STEP,
        "lon_step": lon_step,
        "bounds": {
            "south": round(i * LAT_STEP, 5),
            "north": round((i + 1) * LAT_STEP, 5),
            "west": round(j * lon_step, 5),
            "east": round((j + 1) * lon_step, 5),
        },
        "area_sq_mi": round(height_mi * width_mi, 2),
    }


def snap(lat: float, lon: float) -> dict:
    """Return the grid cell (id + center + span) containing (lat, lon)."""
    lat = min(max(lat, -90.0), 90.0)
    lon = ((lon + 180.0) % 360.0) - 180.0  # wrap into [-180, 180)
    i = math.floor(lat / LAT_STEP)
    j = math.floor(lon / _lon_step((i + 0.5) * LAT_STEP))
    return _cell(i, j)


def neighbors(cell: dict) -> list[dict]:
    """The up-to-8 cells surrounding one cell. Steps one cell width from the
    center and re-snaps, so adjacent rows — whose longitude step differs —
    resolve to whichever cell actually contains the stepped point. Rows past
    the poles are skipped; longitude wraps across the antimeridian (both via
    snap). Deduplicated (near the poles steps can collapse onto one cell)."""
    out: dict[str, dict] = {}
    for di in (-1, 0, 1):
        for dj in (-1, 0, 1):
            if not di and not dj:
                continue
            lat = cell["center_lat"] + di * LAT_STEP
            if abs(lat) > 90.0:
                continue
            n = snap(lat, cell["center_lon"] + dj * cell["lon_step"])
            if n["id"] != cell["id"]:
                out[n["id"]] = n
    return list(out.values())


def from_id(cell_id: str) -> dict:
    """Rebuild the full cell dict from a cache id ("i_j") — the inverse of snap().
    Lets offline tooling (the migrate script) recover a cell from its parquet
    filename alone. Raises ValueError on a malformed id."""
    i, j = (int(p) for p in cell_id.split("_"))
    return _cell(i, j)