Encoding H3 Hexagon Cells as Model Features

from __future__ import annotations

import logging
from dataclasses import dataclass

import geopandas as gpd
import h3
import numpy as np
import pandas as pd

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class H3Config:
    """Pinned H3 settings. Ship this with the model artifact."""
    base_res: int = 8            # finest resolution used as a key
    parent_levels: tuple = (7, 6)
    ring_k: int = 1              # neighbourhood radius in cells

    def validate(self) -> None:
        if not 0 <= self.base_res <= 15:
            raise ValueError(f"base_res {self.base_res} outside H3 range 0-15")
        if any(p >= self.base_res for p in self.parent_levels):
            raise ValueError("parent levels must be coarser than base_res")


def add_h3_keys(gdf: gpd.GeoDataFrame, cfg: H3Config) -> gpd.GeoDataFrame:
    """Attach H3 keys at the base resolution and each configured parent level.

    Requires point geometries in EPSG:4326, because H3 indexes latitude/longitude.
    """
    cfg.validate()
    if gdf.crs is None:
        raise ValueError("GeoDataFrame has no CRS")
    pts = gdf.to_crs("EPSG:4326") if gdf.crs.to_epsg() != 4326 else gdf
    if not (pts.geometry.geom_type == "Point").all():
        pts = pts.copy()
        pts["geometry"] = pts.geometry.centroid

    out = gdf.copy()
    out[f"h3_{cfg.base_res}"] = [
        h3.latlng_to_cell(y, x, cfg.base_res)
        for x, y in zip(pts.geometry.x, pts.geometry.y)
    ]
    for level in cfg.parent_levels:
        out[f"h3_{level}"] = [h3.cell_to_parent(c, level) for c in out[f"h3_{cfg.base_res}"]]

    counts = out[f"h3_{cfg.base_res}"].value_counts()
    logger.info("res %d: %d cells, median %d samples/cell, %d singleton cells",
                cfg.base_res, len(counts), int(counts.median()), int((counts == 1).sum()))
    return out


def ring_aggregate(frame: pd.DataFrame, cfg: H3Config, value_col: str,
                   stat: str = "mean") -> pd.Series:
    """Aggregate a value over each cell's k-ring neighbourhood.

    Uses only the cells present in `frame`; missing neighbours are simply absent
    from the aggregate rather than imputed, which keeps edge cells honest.
    """
    key = f"h3_{cfg.base_res}"
    per_cell = frame.groupby(key)[value_col].agg(stat)
    lookup = per_cell.to_dict()

    values = []
    for cell in frame[key]:
        neighbours = h3.grid_disk(cell, cfg.ring_k)
        present = [lookup[n] for n in neighbours if n in lookup]
        values.append(float(np.mean(present)) if present else np.nan)
    return pd.Series(values, index=frame.index, name=f"{value_col}_ring{cfg.ring_k}")


def cell_area_km2(cell: str) -> float:
    """Real area of a specific cell — hexagons are not all the same size."""
    return h3.cell_area(cell, unit="km^2")

cfg = H3Config(base_res=8, parent_levels=(7, 6), ring_k=1)
samples = add_h3_keys(gpd.read_file("observations.gpkg"), cfg)
print(samples[["h3_8", "h3_7", "h3_6"]].head())

for res in range(5, 11):
    keys = [h3.latlng_to_cell(y, x, res)
            for x, y in zip(samples.geometry.x, samples.geometry.y)]
    n_cells = len(set(keys))
    print(f"res {res:>2}: {n_cells:>7,} cells, {len(samples) / n_cells:>7.1f} samples/cell")

folds = samples["h3_6"].astype("category").cat.codes.to_numpy() % 5   # coarse spatial folds

encoder = SmoothedTargetEncoder(smoothing=20.0)
samples["h3_8_te"] = encoder.fit_transform_oof(
    samples, key="h3_8", target="yield_t_ha", folds=folds
)

samples["ndvi_ring1"] = ring_aggregate(samples, cfg, value_col="ndvi_mean")
samples["ndvi_delta"] = samples["ndvi_mean"] - samples["ndvi_ring1"]
print(samples[["ndvi_mean", "ndvi_ring1", "ndvi_delta"]].describe())

import json

with open("artifacts/h3_config.json", "w") as fh:
    json.dump({"base_res": cfg.base_res, "parent_levels": list(cfg.parent_levels),
               "ring_k": cfg.ring_k}, fh)
encoder.save("artifacts/h3_target_encoder.json")

import h3
import numpy as np


def test_index_is_not_ordinal():
    """Numeric ordering of H3 indices must not track geographic proximity."""
    a = h3.latlng_to_cell(52.5200, 13.4050, 8)     # Berlin
    b = h3.latlng_to_cell(52.5210, 13.4060, 8)     # ~130 m away
    c = h3.latlng_to_cell(-33.8688, 151.2093, 8)   # Sydney
    near = abs(int(a, 16) - int(b, 16))
    far = abs(int(a, 16) - int(c, 16))
    assert not (near < far and near < 10), (
        "if this ever passes as an ordering, the test is wrong — indices are not distances"
    )


def test_parent_contains_child():
    cell = h3.latlng_to_cell(52.52, 13.405, 9)
    parent = h3.cell_to_parent(cell, 6)
    assert cell in h3.cell_to_children(parent, 9)


def test_ring_is_symmetric():
    """Every neighbour of a cell must list that cell as its own neighbour."""
    cell = h3.latlng_to_cell(52.52, 13.405, 8)
    for n in h3.grid_disk(cell, 1):
        assert cell in h3.grid_disk(n, 1)


def test_resolution_pinned_between_train_and_infer():
    cfg = H3Config(base_res=8)
    saved = {"base_res": 8}
    assert cfg.base_res == saved["base_res"], "resolution drifted between train and serve"