Source code: AI habitat classification shared module

tutorial-code/
├── ai-habitat-classification-common/   <- this folder
│   ├── habitat_common.py
│   ├── prepare_data.py
│   ├── requirements.txt
│   ├── README.md
│   ├── samples/
│   │   ├── habitat_samples.csv
│   │   └── blocks.csv
│   └── data/                           <- download cache, created on first run
├── habitat-mapping-sentinel-2-python/   <- Part 1 scripts
├── habitat-classification-lightgbm-sentinel/
└── ...

"""Shared data foundation for the "AI habitat classification with Python" series (Peak District, UK).

Every part of the series imports this module so that all parts use the same study area, 10 m grid,
class scheme, labels, Sentinel-2/Sentinel-1 composites and spatial train/test split.

    import sys; sys.path.insert(0, "../ai-habitat-classification-common")
    import habitat_common as hc

See README.md in this folder for the full interface.
"""
from __future__ import annotations

import json
import time
import warnings
from pathlib import Path

import numpy as np
import pandas as pd

HERE = Path(__file__).resolve().parent
DATA = HERE / "data"              # downloads and cached composites (gitignored)
SAMPLES_CSV = HERE / "samples" / "habitat_samples.csv"   # committed: the shared sample design
BLOCKS_CSV = HERE / "samples" / "blocks.csv"             # committed: block_id -> split and fold
LABELS_GPKG = DATA / "living_england_aoi.gpkg"
LABELS_TIF = DATA / "labels_10m.tif"

# --- Study area and grid ----------------------------------------------------------------------------
CRS = "EPSG:27700"                               # British National Grid
RES = 10                                         # metres
AOI_BOUNDS = (400000, 377500, 425000, 400000)    # xmin, ymin, xmax, ymax: Dark Peak, Edale and Hope valleys
AOI_NAME = "Dark Peak (Kinder Scout, Bleaklow, Edale and the Hope Valley)"
YEAR = 2023                                      # imagery year: matches the Living England 2022-23 epoch
BLOCK_SIZE = 2500                                # spatial blocks for the train/test split, metres
BLOCK_BUFFER = 50                                # no samples within this distance of a block edge, metres
SEED = 42

# --- Class scheme -----------------------------------------------------------------------------------
# id: value in label rasters and samples (0 = no label). le: Living England 2022-23 "Prmry_H" values.
CLASSES = [
    {"id": 1, "key": "bog", "name": "Blanket bog", "colour": "#7b3294",
     "le": ["Bog"]},
    {"id": 2, "key": "heath", "name": "Heath", "colour": "#c2a5cf",
     "le": ["Dwarf Shrub Heath"]},
    {"id": 3, "key": "bracken", "name": "Bracken", "colour": "#e08214",
     "le": ["Bracken"]},
    {"id": 4, "key": "seminat_grass", "name": "Semi-natural grassland", "colour": "#fee08b",
     "le": ["Unimproved Grassland"]},             # acid on the gritstone, calcareous on the limestone near Castleton
    {"id": 5, "key": "improved_grass", "name": "Improved grassland", "colour": "#a6d96a",
     "le": ["Improved and Semi-Improved Grassland", "Arable and Horticultural"]},
    {"id": 6, "key": "broadleaf", "name": "Broadleaf woodland", "colour": "#1a9850",
     "le": ["Broadleaved, Mixed and Yew Woodland", "Scrub"]},
    {"id": 7, "key": "conifer", "name": "Conifer woodland", "colour": "#00441b",
     "le": ["Coniferous Woodland"]},
    {"id": 8, "key": "water", "name": "Water", "colour": "#2166ac",
     "le": ["Water"]},
    {"id": 9, "key": "built_bare", "name": "Built-up and bare", "colour": "#636363",
     "le": ["Built-up Areas and Gardens", "Bare Ground", "Bare Sand"]},
]
EXCLUDED_LE = ["Fen, Marsh and Swamp"]           # too rare and ambiguous here: left unlabelled
CLASS_IDS = [c["id"] for c in CLASSES]
CLASS_NAMES = {c["id"]: c["name"] for c in CLASSES}
CLASS_KEYS = {c["id"]: c["key"] for c in CLASSES}
LE_TO_ID = {le: c["id"] for c in CLASSES for le in c["le"]}
RELIABLE = ["High", "Very High"]                 # Living England reliability values used for sampling

# --- Sentinel-2 / Sentinel-1 -----------------------------------------------------------------------
STAC_URL = "https://planetarycomputer.microsoft.com/api/stac/v1"
S2_BANDS = ["B02", "B03", "B04", "B05", "B06", "B07", "B08", "B8A", "B11", "B12"]
SCL_CLEAR = [4, 5, 6, 7]          # vegetation, bare soil, water, unclassified
SCL_CLOUD = [3, 8, 9, 10]         # cloud shadow, cloud medium/high probability, thin cirrus (dilated by 20 m)
SEASONS = {"winter": [1, 2, 12], "spring": [3, 4, 5], "summer": [6, 7, 8], "autumn": [9, 10, 11]}

LE_SERVICE = ("https://services.arcgis.com/JJzESW51TqeY9uat/ArcGIS/rest/services/"
              "Living_England_2022_2023/FeatureServer")
LE_LAYERS = [3, 5, 6]             # biogeographic zones that intersect the AOI


def geobox():
    """The shared 10 m GeoBox (odc-geo) in EPSG:27700: 2,500 columns x 2,250 rows."""
    from odc.geo.geobox import GeoBox
    return GeoBox.from_bbox(AOI_BOUNDS, crs=CRS, resolution=RES)


def transform():
    """Affine transform of the shared grid (rasterio convention, north-up)."""
    from rasterio.transform import from_origin
    return from_origin(AOI_BOUNDS[0], AOI_BOUNDS[3], RES, RES)


def shape():
    """(rows, cols) of the shared grid."""
    return (int((AOI_BOUNDS[3] - AOI_BOUNDS[1]) / RES), int((AOI_BOUNDS[2] - AOI_BOUNDS[0]) / RES))


def aoi_bbox_wgs84():
    """AOI bounding box in longitude/latitude, for STAC searches."""
    from pyproj import Transformer
    t = Transformer.from_crs(CRS, "EPSG:4326", always_xy=True)
    xs, ys = t.transform([AOI_BOUNDS[0], AOI_BOUNDS[0], AOI_BOUNDS[2], AOI_BOUNDS[2]],
                         [AOI_BOUNDS[1], AOI_BOUNDS[3], AOI_BOUNDS[1], AOI_BOUNDS[3]])
    return (min(xs), min(ys), max(xs), max(ys))


def class_cmap():
    """(ListedColormap, BoundaryNorm) for label/prediction rasters with values 1..9 (mask 0 before plotting)."""
    from matplotlib.colors import BoundaryNorm, ListedColormap
    cmap = ListedColormap([c["colour"] for c in CLASSES])
    cmap.set_bad(alpha=0)
    return cmap, BoundaryNorm(np.arange(0.5, len(CLASSES) + 1.5), cmap.N)


# --- Labels ------------------------------------------------------------------------------------------
def _get_json(url, params, tries=5):
    import requests
    for i in range(tries):
        try:
            r = requests.post(url, data=params, timeout=180)   # POST: long objectIds lists overflow a GET URL
            r.raise_for_status()
            js = r.json()
            if "error" in js:
                raise RuntimeError(js["error"])
            return js
        except Exception:
            if i == tries - 1:
                raise
            time.sleep(5 * (i + 1))


def download_living_england(force=False, verbose=True):
    """Download Living England 2022-23 polygons that intersect the AOI (Natural England, OGL v3).

    Pages through the public ArcGIS FeatureServer (no login), clips to the AOI and saves a GeoPackage.
    Returns a GeoDataFrame in EPSG:27700 with columns le_segid, le_habitat, le_reliability, le_source, class_id.
    """
    import geopandas as gpd
    from shapely.geometry import box

    if LABELS_GPKG.exists() and not force:
        return gpd.read_file(LABELS_GPKG)
    DATA.mkdir(exist_ok=True)
    env = ",".join(str(v) for v in AOI_BOUNDS)
    frames = []
    for layer in LE_LAYERS:
        url = f"{LE_SERVICE}/{layer}/query"
        ids = _get_json(url, {"where": "1=1", "geometry": env, "geometryType": "esriGeometryEnvelope",
                              "inSR": 27700, "spatialRel": "esriSpatialRelIntersects",
                              "returnIdsOnly": "true", "f": "json"})["objectIds"] or []
        ids = sorted(ids)
        for i in range(0, len(ids), 1000):
            js = _get_json(url, {"objectIds": ",".join(map(str, ids[i:i + 1000])),
                                 "outFields": "SegID,Prmry_H,Relblty,Source", "outSR": 27700, "f": "geojson"})
            frames.append(gpd.GeoDataFrame.from_features(js["features"], crs=CRS))
        if verbose:
            print(f"Living England 2022-23, zone layer {layer}: {len(ids):,} polygons intersect the AOI")
    gdf = gpd.GeoDataFrame(pd.concat(frames, ignore_index=True), geometry="geometry", crs=CRS)
    gdf = gdf.rename(columns={"Prmry_H": "le_habitat", "Relblty": "le_reliability", "Source": "le_source",
                              "SegID": "le_segid"})
    gdf["geometry"] = gdf.geometry.make_valid().intersection(box(*AOI_BOUNDS))
    gdf = gdf[~gdf.geometry.is_empty].explode(index_parts=False)
    gdf = gdf[gdf.geometry.geom_type == "Polygon"]
    gdf["class_id"] = gdf["le_habitat"].map(LE_TO_ID).fillna(0).astype("uint8")
    gdf = gdf[["le_segid", "le_habitat", "le_reliability", "le_source", "class_id", "geometry"]].reset_index(drop=True)
    gdf.to_file(LABELS_GPKG, driver="GPKG")
    return gdf


def load_labels():
    """Living England polygons for the AOI (downloads on first use)."""
    return download_living_england()


def rasterize_labels(gdf=None, reliable_only=False):
    """Rasterise class_id onto the shared grid (uint8, 0 = unlabelled). Pixel-centre rule."""
    from rasterio.features import rasterize
    gdf = load_labels() if gdf is None else gdf
    if reliable_only:
        gdf = gdf[gdf["le_reliability"].isin(RELIABLE)]
    gdf = gdf[gdf["class_id"] > 0]
    return rasterize(zip(gdf.geometry, gdf["class_id"]), out_shape=shape(), transform=transform(), fill=0,
                     dtype="uint8")


def label_raster():
    """Full Living England label map on the shared grid (all reliabilities), cached as data/labels_10m.tif."""
    import rasterio
    if LABELS_TIF.exists():
        with rasterio.open(LABELS_TIF) as src:
            return src.read(1)
    arr = rasterize_labels()
    write_geotiff(LABELS_TIF, arr, nodata=0)
    return arr


def write_geotiff(path, arr, nodata=None, dtype=None):
    """Write a 2-D (rows, cols) or 3-D (bands, rows, cols) array on the shared grid as a compressed GeoTIFF."""
    import rasterio
    arr = np.asarray(arr)
    if arr.ndim == 2:
        arr = arr[None]
    dtype = dtype or arr.dtype
    Path(path).parent.mkdir(parents=True, exist_ok=True)
    with rasterio.open(path, "w", driver="GTiff", height=arr.shape[1], width=arr.shape[2], count=arr.shape[0],
                       dtype=dtype, crs=CRS, transform=transform(), nodata=nodata, compress="deflate",
                       tiled=True, blockxsize=512, blockysize=512) as dst:
        dst.write(arr.astype(dtype))


# --- Samples and the spatial split -------------------------------------------------------------------
def block_id(x, y):
    """Spatial block index (row-major from the north-west corner) for BNG coordinates."""
    ncols = int(np.ceil((AOI_BOUNDS[2] - AOI_BOUNDS[0]) / BLOCK_SIZE))
    bc = ((np.asarray(x) - AOI_BOUNDS[0]) // BLOCK_SIZE).astype(int)
    br = ((AOI_BOUNDS[3] - np.asarray(y)) // BLOCK_SIZE).astype(int)
    return br * ncols + bc


def load_blocks():
    """GeoDataFrame of the spatial blocks: block_id, split ("train"/"test"), fold, geometry (samples/blocks.csv)."""
    import geopandas as gpd
    from shapely.geometry import box
    df = pd.read_csv(BLOCKS_CSV)
    ncols = int(np.ceil((AOI_BOUNDS[2] - AOI_BOUNDS[0]) / BLOCK_SIZE))
    geoms = []
    for b in df["block_id"]:
        r, c = divmod(int(b), ncols)
        x0, y1 = AOI_BOUNDS[0] + c * BLOCK_SIZE, AOI_BOUNDS[3] - r * BLOCK_SIZE
        geoms.append(box(x0, max(y1 - BLOCK_SIZE, AOI_BOUNDS[1]), min(x0 + BLOCK_SIZE, AOI_BOUNDS[2]), y1))
    return gpd.GeoDataFrame(df, geometry=geoms, crs=CRS)


def load_samples(split=None):
    """The committed sample design (samples/habitat_samples.csv) as a DataFrame.

    Columns: sample_id, x, y (BNG pixel centre), row, col (indices into the shared grid), class_id, class_key,
    le_habitat, le_reliability, le_source, block_id, split ("train" or "test"), fold (0-4 for train, -1 for test).
    """
    df = pd.read_csv(SAMPLES_CSV)
    return df if split is None else df[df["split"] == split].reset_index(drop=True)


def sample_values(arr, samples):
    """Values of a (rows, cols) or (bands, rows, cols) array at each sample's row/col."""
    arr = np.asarray(arr)
    r, c = samples["row"].to_numpy(), samples["col"].to_numpy()
    return arr[r, c] if arr.ndim == 2 else arr[:, r, c].T


# --- Sentinel-2 --------------------------------------------------------------------------------------
def stac_client():
    import planetary_computer
    import pystac_client
    return pystac_client.Client.open(STAC_URL, modifier=planetary_computer.sign_inplace)


def search_s2(start=f"{YEAR}-01-01", end=f"{YEAR}-12-31", max_cloud=90):
    """Sentinel-2 L2A items over the AOI, one per date and tile (newest processing baseline kept)."""
    items = stac_client().search(collections=["sentinel-2-l2a"], bbox=aoi_bbox_wgs84(), datetime=f"{start}/{end}",
                                 query={"eo:cloud_cover": {"lt": max_cloud}}).item_collection()
    best = {}
    for it in items:
        key = (it.datetime.date(), it.properties["s2:mgrs_tile"])
        if key not in best or it.properties["s2:processing_baseline"] > best[key].properties["s2:processing_baseline"]:
            best[key] = it
    return sorted(best.values(), key=lambda i: i.datetime)


def load_s2(items, bands=S2_BANDS, chunks=None):
    """Lazy Sentinel-2 cube on the shared grid: dask-backed xarray.Dataset of raw uint16 DNs plus SCL.

    Reflectance = (DN - offset) / 10000, with the per-date offset from s2_offsets(). 20 m bands are
    resampled bilinearly to the 10 m grid; SCL with nearest neighbour.
    """
    import odc.stac
    odc.stac.configure_rio(cloud_defaults=True)
    ds = odc.stac.load(items, bands=list(bands) + ["SCL"], geobox=geobox(), groupby="solar_day",
                       chunks=chunks or {"time": 1, "x": 2500, "y": 2250}, dtype="uint16", nodata=0,
                       resampling={"*": "bilinear", "SCL": "nearest"})
    return ds.assign_coords(time=ds.time.dt.floor("D"))


def s2_offsets(items):
    """Per-date BOA offset (1000 for processing baseline >= 04.00, else 0) as a pandas Series."""
    return pd.Series({pd.Timestamp(it.datetime.date()):
                      1000 if float(it.properties["s2:processing_baseline"]) >= 4.0 else 0 for it in items})


def s2_clear_masks(items=None, cache=True):
    """Clear-sky masks for every date: (bool DataArray time x y x x, clear-fraction Series).

    clear = SCL in SCL_CLEAR and not within 20 m of SCL_CLOUD. Cached as data/s2_clear_<year>.npz.
    """
    import xarray as xr
    from scipy import ndimage
    items = items or search_s2()
    f = DATA / f"s2_clear_{items[0].datetime.year}.npz"
    if cache and f.exists():
        z = np.load(f)
        clear = xr.DataArray(np.unpackbits(z["bits"], axis=0, count=z["times"].size).astype(bool),
                             dims=("time", "y", "x"), coords={"time": pd.to_datetime(z["times"])}, name="clear")
    else:
        scl = load_s2(items, bands=[])["SCL"].compute()
        cloudy = np.isin(scl.values, SCL_CLOUD)
        grown = ndimage.binary_dilation(cloudy, iterations=2, axes=(1, 2))
        clear = xr.DataArray(np.isin(scl.values, SCL_CLEAR) & ~grown, dims=("time", "y", "x"),
                             coords={"time": scl.time.values}, name="clear")
        DATA.mkdir(exist_ok=True)
        np.savez_compressed(f, bits=np.packbits(clear.values, axis=0), times=clear.time.values.astype("datetime64[D]"))
    frac = clear.mean(("y", "x")).to_series()
    return clear, frac


def s2_composite(name, months, year=YEAR, bands=S2_BANDS, min_clear=0.05, cache=True, verbose=True):
    """Cloud-masked median surface reflectance for the given months of one year on the shared grid.

    Returns an xarray.Dataset of float32 reflectance (0-1), one variable per band, NaN where no clear
    view exists, plus "n_clear" (number of clear dates per pixel). Cached as data/s2_<year>_<name>.tif
    (int16 reflectance x 10000, nodata -9999; last band n_clear).
    """
    import rasterio
    import xarray as xr
    f = DATA / f"s2_{year}_{name}.tif"
    if cache and f.exists():
        with rasterio.open(f) as src:
            arr = src.read().astype("float32")
        ds = xr.Dataset({b: (("y", "x"), np.where(arr[i] == -9999, np.nan, arr[i] / 10000))
                         for i, b in enumerate(bands)})
        ds["n_clear"] = (("y", "x"), arr[-1])
        return ds
    import dask
    items = search_s2(f"{year}-01-01", f"{year}-12-31")
    clear, frac = s2_clear_masks(items, cache=cache)
    times = frac.index[(frac >= min_clear) & frac.index.month.isin(months)]
    t0 = time.time()
    ds = load_s2(items, bands=bands).sel(time=times)
    off = xr.DataArray(s2_offsets(items).reindex(times).to_numpy(), dims="time", coords={"time": ds.time})
    refl = ((ds[list(bands)].astype("float32") - off) / 10000).where(clear.sel(time=times).values)
    with dask.config.set(scheduler="threads", num_workers=16), warnings.catch_warnings():
        warnings.simplefilter("ignore", RuntimeWarning)
        med = refl.chunk({"time": -1, "y": 450, "x": 500}).median("time").compute()
    n = clear.sel(time=times).sum("time").values.astype("float32")
    stack = [np.where(np.isnan(med[b].values), -9999, np.round(med[b].values * 10000)) for b in bands]
    write_geotiff(f, np.stack(stack + [n]), nodata=-9999, dtype="int16")
    if verbose:
        print(f"S2 {name} {year}: {len(times)} dates with >= {min_clear:.0%} clear sky, built in {time.time() - t0:.0f} s")
    out = xr.Dataset({b: (("y", "x"), med[b].values.astype("float32")) for b in bands})
    out["n_clear"] = (("y", "x"), n)
    return out


def s2_seasonal_composites(year=YEAR, **kw):
    """Dict season -> s2_composite(...) for SEASONS (winter = Jan, Feb and Dec of the same year)."""
    return {s: s2_composite(s, m, year=year, **kw) for s, m in SEASONS.items()}


def s2_dates(year=YEAR, min_clear=0.05):
    """Clear fraction of the AOI for each date in the year with at least min_clear clear sky (pandas Series)."""
    items = search_s2(f"{year}-01-01", f"{year}-12-31")
    _, frac = s2_clear_masks(items)
    return frac[frac >= min_clear]


def s2_sample_timeseries(samples=None, year=YEAR, bands=S2_BANDS, min_clear=0.05, cache=True, verbose=True):
    """Per-date Sentinel-2 reflectance at every sample pixel (for time-series models).

    Returns a long DataFrame: sample_id, date, clear (bool), then one float32 reflectance column per band
    (NaN where not clear). Covers every date with >= min_clear of the AOI clear (50 dates in 2023).
    Cached as data/s2_<year>_samples_timeseries.parquet; built for the full committed sample set.
    """
    f = DATA / f"s2_{year}_samples_timeseries.parquet"
    if cache and f.exists():
        ts = pd.read_parquet(f)
    else:
        allsamp = load_samples()
        items = search_s2(f"{year}-01-01", f"{year}-12-31")
        clear, frac = s2_clear_masks(items, cache=cache)
        times = frac.index[frac >= min_clear]
        off = s2_offsets(items)
        ds = load_s2(items, bands=bands).sel(time=times)
        r, c = allsamp["row"].to_numpy(), allsamp["col"].to_numpy()
        t0, parts = time.time(), []
        for i, t in enumerate(times):
            day = ds.sel(time=t).compute()
            ok = clear.sel(time=t).values[r, c]
            part = pd.DataFrame({"sample_id": allsamp["sample_id"].to_numpy(), "date": t, "clear": ok})
            for b in bands:
                v = (day[b].values[r, c].astype("float32") - off[t]) / 10000
                part[b] = np.where(ok, v, np.nan).astype("float32")
            parts.append(part)
            if verbose:
                print(f"  S2 samples {i + 1}/{len(times)} {t.date()} ({time.time() - t0:.0f} s)", flush=True)
        ts = pd.concat(parts, ignore_index=True)
        DATA.mkdir(exist_ok=True)
        ts.to_parquet(f, index=False)
    if samples is not None:
        ts = ts[ts["sample_id"].isin(samples["sample_id"])].reset_index(drop=True)
    return ts


# --- Sentinel-1 RTC ----------------------------------------------------------------------------------
def search_s1(start=f"{YEAR}-01-01", end=f"{YEAR}-12-31", orbit_state=None):
    """Sentinel-1 RTC (gamma0, radiometrically terrain corrected) items over the AOI."""
    q = {"sat:orbit_state": {"eq": orbit_state}} if orbit_state else None
    items = stac_client().search(collections=["sentinel-1-rtc"], bbox=aoi_bbox_wgs84(), datetime=f"{start}/{end}",
                                 query=q).item_collection()
    return sorted(items, key=lambda i: i.datetime)


def load_s1(items, chunks=None):
    """Lazy Sentinel-1 RTC cube (vv, vh, linear gamma0, float32) on the shared grid, grouped by solar day."""
    import odc.stac
    odc.stac.configure_rio(cloud_defaults=True)
    ds = odc.stac.load(items, bands=["vv", "vh"], geobox=geobox(), groupby="solar_day",
                       chunks=chunks or {"time": 1, "x": 2500, "y": 2250}, dtype="float32", nodata=np.nan,
                       resampling="bilinear")
    return ds.assign_coords(time=ds.time.dt.floor("D"))


def s1_composite(name, months, year=YEAR, orbit_state=None, cache=True, verbose=True):
    """Median VV and VH gamma0 for the given months of one year on the shared grid.

    Returns an xarray.Dataset with vv, vh (linear), vv_db, vh_db and ratio_db (VH - VV in dB).
    Cached as data/s1_<year>_<name>[_<orbit_state>].tif (float32 linear vv, vh).
    """
    import rasterio
    import xarray as xr
    f = DATA / f"s1_{year}_{name}{'_' + orbit_state if orbit_state else ''}.tif"
    if cache and f.exists():
        with rasterio.open(f) as src:
            vv, vh = src.read().astype("float32")
    else:
        import dask
        t0 = time.time()
        items = [i for i in search_s1(f"{year}-01-01", f"{year}-12-31", orbit_state) if i.datetime.month in months]
        ds = load_s1(items)
        with dask.config.set(scheduler="threads", num_workers=16), warnings.catch_warnings():
            warnings.simplefilter("ignore", RuntimeWarning)
            med = ds.where(ds > 0).chunk({"time": -1, "y": 450, "x": 500}).median("time").compute()
        vv, vh = med["vv"].values.astype("float32"), med["vh"].values.astype("float32")
        write_geotiff(f, np.stack([vv, vh]), nodata=np.nan, dtype="float32")
        if verbose:
            print(f"S1 {name} {year}: {ds.sizes['time']} dates, built in {time.time() - t0:.0f} s")
    with np.errstate(divide="ignore", invalid="ignore"):
        out = xr.Dataset({"vv": (("y", "x"), vv), "vh": (("y", "x"), vh),
                          "vv_db": (("y", "x"), 10 * np.log10(vv)), "vh_db": (("y", "x"), 10 * np.log10(vh))})
    out["ratio_db"] = out["vh_db"] - out["vv_db"]
    return out


def describe():
    """A short summary of the shared configuration."""
    return {"crs": CRS, "res": RES, "aoi_bounds": AOI_BOUNDS, "shape": shape(), "year": YEAR,
            "block_size": BLOCK_SIZE, "classes": {c["id"]: c["name"] for c in CLASSES}}


if __name__ == "__main__":
    print(json.dumps(describe(), indent=2))

"""Prepare the shared labels and sample design for the "AI habitat classification with Python" series.

1. Downloads Living England 2022-23 habitat polygons for the Dark Peak AOI (Natural England, OGL v3, no login).
2. Maps them to the 9-class scheme in habitat_common.CLASSES and rasterises them onto the shared 10 m grid.
3. Draws a stratified random sample of "pure" pixels from high-reliability polygons and assigns each to a
   2.5 km spatial block; blocks are split 70/30 into train/test with a fixed seed, and train blocks into
   5 folds for blocked cross-validation.
4. Writes samples/habitat_samples.csv and samples/blocks.csv (committed, so every part uses the same split).

Run:  python prepare_data.py               labels + samples (samples are only redrawn with --resample)
      python prepare_data.py --composites  also build the 2023 seasonal Sentinel-2 and Sentinel-1 composites
      python prepare_data.py --timeseries  also extract per-date Sentinel-2 reflectance at every sample (Part 4)
"""
import argparse
import hashlib

import numpy as np
import pandas as pd
from scipy import ndimage

import habitat_common as hc

N_PER_CLASS = 1000      # samples per class (fewer if a class has fewer eligible pixels)
TEST_FRACTION = 0.3     # share of blocks held out for testing
N_FOLDS = 5


def draw_samples(labels_all, labels_rel):
    rng = np.random.default_rng(hc.SEED)
    rows, cols = hc.shape()
    rr, cc = np.mgrid[0:rows, 0:cols]
    x = hc.AOI_BOUNDS[0] + (cc + 0.5) * hc.RES
    y = hc.AOI_BOUNDS[3] - (rr + 0.5) * hc.RES
    # keep pixel centres at least BLOCK_BUFFER from any block edge, so test pixels never touch training pixels
    off_x = (x - hc.AOI_BOUNDS[0]) % hc.BLOCK_SIZE
    off_y = (hc.AOI_BOUNDS[3] - y) % hc.BLOCK_SIZE
    inner = ((off_x > hc.BLOCK_BUFFER) & (off_x < hc.BLOCK_SIZE - hc.BLOCK_BUFFER)
             & (off_y > hc.BLOCK_BUFFER) & (off_y < hc.BLOCK_SIZE - hc.BLOCK_BUFFER))
    out = []
    print(f"\n{'class':<20}{'eligible px':>12}{'sampled':>9}")
    for c in hc.CLASSES:
        k = c["id"]
        # "pure" pixels: reliable label k, and all 8 neighbours also have label k in the full map
        pure = ndimage.binary_erosion((labels_rel == k) & (labels_all == k), structure=np.ones((3, 3)),
                                      border_value=0)
        idx = np.flatnonzero(pure & inner)
        pick = np.sort(rng.choice(idx, size=min(N_PER_CLASS, idx.size), replace=False))
        r, cl = np.unravel_index(pick, (rows, cols))
        out.append(pd.DataFrame({"x": x[r, cl], "y": y[r, cl], "row": r, "col": cl, "class_id": k,
                                 "class_key": c["key"]}))
        print(f"{c['name']:<20}{idx.size:>12,}{pick.size:>9,}")
    return pd.concat(out, ignore_index=True)


def assign_blocks(samples):
    rng = np.random.default_rng(hc.SEED)
    ncols = int(np.ceil((hc.AOI_BOUNDS[2] - hc.AOI_BOUNDS[0]) / hc.BLOCK_SIZE))
    nrows = int(np.ceil((hc.AOI_BOUNDS[3] - hc.AOI_BOUNDS[1]) / hc.BLOCK_SIZE))
    blocks = np.arange(ncols * nrows)
    order = rng.permutation(blocks)
    n_test = int(round(TEST_FRACTION * blocks.size))
    test = set(order[:n_test].tolist())
    train_blocks = [b for b in order if b not in test]
    fold = {b: i % N_FOLDS for i, b in enumerate(train_blocks)}
    bdf = pd.DataFrame({"block_id": blocks, "split": ["test" if b in test else "train" for b in blocks],
                        "fold": [-1 if b in test else fold[b] for b in blocks]})
    samples["block_id"] = hc.block_id(samples["x"], samples["y"])
    samples = samples.merge(bdf, on="block_id", how="left")
    return samples, bdf


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--resample", action="store_true", help="redraw samples/*.csv even if they exist")
    ap.add_argument("--composites", action="store_true", help="also build the S2 and S1 seasonal composites")
    ap.add_argument("--timeseries", action="store_true", help="also extract per-date S2 reflectance at every sample")
    args = ap.parse_args()

    print("=== Living England 2022-23 labels ===")
    gdf = hc.load_labels()
    print(f"{len(gdf):,} polygons clipped to the AOI {hc.AOI_BOUNDS} ({hc.CRS})")
    area = gdf.assign(km2=gdf.area / 1e6).groupby(["le_habitat"])["km2"].sum().sort_values(ascending=False)
    rel = gdf.assign(km2=gdf.area / 1e6).groupby("le_reliability")["km2"].sum()
    print(area.round(2).to_string())
    print("\narea by reliability (km2):\n" + rel.round(1).to_string())

    labels_all = hc.label_raster()
    labels_rel = hc.rasterize_labels(gdf, reliable_only=True)
    print(f"\nshared grid: {hc.shape()[1]} x {hc.shape()[0]} px at {hc.RES} m; "
          f"{(labels_all > 0).mean():.1%} labelled, {(labels_rel > 0).mean():.1%} labelled with High/Very High reliability")
    for c in hc.CLASSES:
        print(f"  {c['id']} {c['name']:<20} {(labels_all == c['id']).mean():6.1%} of the AOI")

    if hc.SAMPLES_CSV.exists() and not args.resample:
        samples = hc.load_samples()
        print(f"\nusing the committed sample design: {hc.SAMPLES_CSV.name} ({len(samples):,} samples)")
    else:
        print("\n=== stratified sample of pure, high-reliability pixels ===")
        samples = draw_samples(labels_all, labels_rel)
        samples, bdf = assign_blocks(samples)
        le = gdf[["le_habitat", "le_reliability", "le_source", "geometry"]]
        import geopandas as gpd
        pts = gpd.GeoDataFrame(samples, geometry=gpd.points_from_xy(samples.x, samples.y), crs=hc.CRS)
        pts = gpd.sjoin(pts, le, how="left", predicate="within").drop_duplicates("row col".split())
        samples = pd.DataFrame(pts.drop(columns=["geometry", "index_right"]))
        samples.insert(0, "sample_id", np.arange(len(samples)))
        cols = ["sample_id", "x", "y", "row", "col", "class_id", "class_key", "le_habitat", "le_reliability",
                "le_source", "block_id", "split", "fold"]
        hc.SAMPLES_CSV.parent.mkdir(exist_ok=True)
        samples[cols].to_csv(hc.SAMPLES_CSV, index=False)
        bdf.to_csv(hc.BLOCKS_CSV, index=False)
        samples = hc.load_samples()

    blocks = pd.read_csv(hc.BLOCKS_CSV)
    print(f"\n=== spatial split: {hc.BLOCK_SIZE / 1000:g} km blocks, seed {hc.SEED} ===")
    print(f"{len(blocks)} blocks: {(blocks.split == 'train').sum()} train, {(blocks.split == 'test').sum()} test")
    tab = pd.crosstab(samples["class_key"], samples["split"]).reindex([c["key"] for c in hc.CLASSES])
    tab["total"] = tab.sum(axis=1)
    print(tab.to_string())
    print(f"total {len(samples):,} samples; {samples.block_id.nunique()} blocks contain samples")
    print(f"sha256(habitat_samples.csv) = {hashlib.sha256(hc.SAMPLES_CSV.read_bytes()).hexdigest()[:16]}")

    if args.composites:
        print("\n=== Sentinel-2 seasonal composites ===")
        frac = hc.s2_dates()
        print(f"{len(frac)} dates in {hc.YEAR} with >= 5% of the AOI clear")
        hc.s2_seasonal_composites()
        print("\n=== Sentinel-1 RTC seasonal composites ===")
        for s, m in hc.SEASONS.items():
            hc.s1_composite(s, m)

    if args.timeseries:
        print("\n=== Sentinel-2 per-date reflectance at every sample ===")
        ts = hc.s2_sample_timeseries()
        print(f"{len(ts):,} rows: {ts.sample_id.nunique():,} samples x {ts.date.nunique()} dates; "
              f"{ts.clear.mean():.0%} of observations clear")


if __name__ == "__main__":
    main()