import logging
import math
from collections import defaultdict

from constants import GRID_CELL_SIZE

log = logging.getLogger("rightmove")


class PostcodeSpatialIndex:
    """Grid-based spatial index over arcgis postcodes for nearest-lookup."""

    def __init__(self, lats: list[float], lngs: list[float], postcodes: list[str]):
        self.grid: dict[tuple[int, int], list[tuple[float, float, str]]] = defaultdict(
            list
        )
        for lat, lng, pcd in zip(lats, lngs, postcodes):
            gx = int(math.floor(lng / GRID_CELL_SIZE))
            gy = int(math.floor(lat / GRID_CELL_SIZE))
            self.grid[(gx, gy)].append((lat, lng, pcd))
        log.info(
            "Postcode spatial index: %d cells, %d postcodes", len(self.grid), len(lats)
        )

    def nearest(self, lat: float, lng: float) -> str | None:
        gx = int(math.floor(lng / GRID_CELL_SIZE))
        gy = int(math.floor(lat / GRID_CELL_SIZE))
        best_dist = float("inf")
        best_pcd = None
        for dx in range(-1, 2):
            for dy in range(-1, 2):
                for plat, plng, pcd in self.grid.get((gx + dx, gy + dy), []):
                    d = (plat - lat) ** 2 + (plng - lng) ** 2
                    if d < best_dist:
                        best_dist = d
                        best_pcd = pcd
        return best_pcd