Blame for preprocess/collect_features.py

"""
RockGlacierNet — Stage 1: feature/mask COLLECTION (Option A: enlarged window).

For every rock-glacier polygon in the inventory shapefile:
  1. compute its bbox in UTM and pad by CONTEXT_BUFFER_M on every side,
  2. download a 12-band Sentinel-2 + terrain feature stack over that window,
  3. rasterize the polygon to the EXACT grid of the downloaded feature tile,
     so feature and mask are co-registered by construction (no harmonize step).

Outputs (and nothing outside patch_dataset/):
  patch_dataset/raw/features/feature_XXXX.tif   (12-band, float32, 10 m, UTM45N)
  patch_dataset/raw/masks/mask_XXXX.tif         (1-band uint8, same grid)

ENVIRONMENT: needs `earthengine-api`, `geemap`, `geopandas`, `rasterio`.
Run on the machine that has them (your Windows box):
    python preprocess/collect_features.py
    python preprocess/collect_features.py --limit 5     # smoke test
    python preprocess/collect_features.py --shapefile /path/to/your.shp

All tunables live in patch_dataset/config.py.
"""

import os
import sys
import argparse

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from patch_dataset import config as C

import numpy as np
import geopandas as gpd
import rasterio
from rasterio.features import rasterize
import ee
import geemap
from shapely.geometry import box


# ─────────────────────────────────────────────────────────────────────────────
# Earth Engine image builders  (band order MUST match C.BAND_NAMES)
# ─────────────────────────────────────────────────────────────────────────────

def get_s2_image(roi):
    coll = (ee.ImageCollection(C.S2_COLLECTION)
            .filterBounds(roi)
            .filterDate(C.S2_DATE_START, C.S2_DATE_END))

    def mask_clouds(image):
        qa = image.select("QA60")
        cloud  = 1 << 10
        cirrus = 1 << 11
        m = qa.bitwiseAnd(cloud).eq(0).And(qa.bitwiseAnd(cirrus).eq(0))
        return image.updateMask(m).divide(C.S2_SR_DIVISOR)

    s2 = coll.map(mask_clouds).median().clip(roi)
    ndvi = s2.normalizedDifference(["B8", "B4"]).rename("NDVI")
    ndwi = s2.normalizedDifference(["B3", "B8"]).rename("NDWI")
    ndsi = s2.normalizedDifference(["B3", "B11"]).rename("NDSI")
    # B2,B3,B4,B8,B11,B12,NDVI,NDWI,NDSI
    return s2.select(["B2", "B3", "B4", "B8", "B11", "B12"]).addBands([ndvi, ndwi, ndsi])


def get_terrain_image(roi):
    # Compute slope/aspect BEFORE mosaicking so the native 30 m grid is preserved.
    def add_terrain(img):
        slope  = ee.Terrain.slope(img).rename("slope")
        aspect = ee.Terrain.aspect(img).rename("aspect")
        return img.addBands([slope, aspect])

    terrain = (ee.ImageCollection(C.DEM_COLLECTION).select("DEM")
               .map(add_terrain).mosaic().clip(roi))
    elevation = terrain.select("DEM").rename("elevation")   # Elev
    slope     = terrain.select("slope")                     # Slope
    aspect    = terrain.select("aspect")                    # Aspect
    return elevation.addBands([slope, aspect])


def build_feature_image(roi):
    """12-band stack in the contractual order C.BAND_NAMES."""
    return get_s2_image(roi).addBands(get_terrain_image(roi))


# ─────────────────────────────────────────────────────────────────────────────
# Per-polygon processing
# ─────────────────────────────────────────────────────────────────────────────

def padded_roi(geom_utm):
    """Return (ee_roi_4326, bounds_utm) for a single UTM geometry + buffer."""
    minx, miny, maxx, maxy = geom_utm.bounds
    b = C.CONTEXT_BUFFER_M
    minx, miny, maxx, maxy = minx - b, miny - b, maxx + b, maxy + b
    box_utm   = gpd.GeoDataFrame(geometry=[box(minx, miny, maxx, maxy)], crs=C.TARGET_CRS)
    box_4326  = box_utm.to_crs("EPSG:4326").total_bounds
    roi = ee.Geometry.BBox(box_4326[0], box_4326[1], box_4326[2], box_4326[3])
    return roi, (minx, miny, maxx, maxy)


def rasterize_aligned_mask(geom_utm, feat_tif_path, out_mask_path):
    """Rasterize geom onto the EXACT grid of the downloaded feature tile."""
    with rasterio.open(feat_tif_path) as src:
        transform, height, width, crs = src.transform, src.height, src.width, src.crs
    mask = rasterize(
        [(geom_utm, 1)],
        out_shape=(height, width),
        transform=transform,
        fill=0,
        dtype="uint8",
    )
    with rasterio.open(
        out_mask_path, "w", driver="GTiff",
        height=height, width=width, count=1, dtype="uint8",
        crs=crs, transform=transform, nodata=0,
    ) as dst:
        dst.write(mask, 1)
    return int(mask.sum())


def main():
    ap = argparse.ArgumentParser(description="Collect enlarged-window feature/mask tiles.")
    ap.add_argument("--shapefile", default=str(C.SHAPEFILE_PATH))
    ap.add_argument("--limit", type=int, default=None, help="process only first N polygons")
    args = ap.parse_args()

    # EE init
    try:
        ee.Initialize(project=C.GEE_PROJECT)
    except Exception:
        ee.Authenticate()
        ee.Initialize(project=C.GEE_PROJECT)

    C.RAW_FEATURES_DIR.mkdir(parents=True, exist_ok=True)
    C.RAW_MASKS_DIR.mkdir(parents=True, exist_ok=True)

    gdf = gpd.read_file(args.shapefile).to_crs(C.TARGET_CRS)
    n = len(gdf) if args.limit is None else min(args.limit, len(gdf))
    print(f"Loaded {len(gdf)} polygons from {args.shapefile} | processing {n} "
          f"| buffer={C.CONTEXT_BUFFER_M:.0f} m | scale={C.SCALE_M:.0f} m")

    n_done = n_skip = n_fail = 0
    for idx in range(n):
        geom = gdf.geometry.iloc[idx]
        if geom is None or geom.is_empty:
            print(f"  [{idx:04d}] empty geometry — skipped")
            n_skip += 1
            continue

        feat_path = C.RAW_FEATURES_DIR / f"feature_{idx:04d}.tif"
        mask_path = C.RAW_MASKS_DIR / f"mask_{idx:04d}.tif"
        if feat_path.exists() and mask_path.exists():
            n_skip += 1
            continue

        try:
            roi, _ = padded_roi(geom)
            img = build_feature_image(roi)
            geemap.ee_export_image(
                img, filename=str(feat_path), scale=C.SCALE_M,
                region=roi, crs=C.TARGET_CRS, file_per_band=False,
            )
            fg = rasterize_aligned_mask(geom, str(feat_path), str(mask_path))
            n_done += 1
            if n_done % 25 == 0 or idx == n - 1:
                print(f"  [{idx:04d}] done={n_done} skip={n_skip} fail={n_fail} "
                      f"(last fg px={fg})")
        except Exception as e:
            n_fail += 1
            print(f"  [{idx:04d}] FAILED: {e}")

    print(f"\nCollection complete. done={n_done} skipped={n_skip} failed={n_fail}")
    print(f"Features -> {C.RAW_FEATURES_DIR}")
    print(f"Masks    -> {C.RAW_MASKS_DIR}")


if __name__ == "__main__":
    main()