"""
RockGlacierNet — Enlarged per-glacier collection ("best of both").  [patch_dataset/]

One crop per inventory polygon (max diversity: ~610 distinct positive examples),
each glacier centered with GLACIER_BUFFER_M of REAL surrounding terrain (not the
zero-padding the old dataset/ relied on). Reuses the bug-fixed download path from
collect_region (memory-light cloud-filtered composite, correct UTM terrain,
MemoryFile-safe download with sub-tiling fallback).

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

Then:
  python patch_dataset/build_patches.py --patch-size 128 \
      --features-dir patch_dataset/raw_glacier/features \
      --masks-dir   patch_dataset/raw_glacier/masks
  python patch_dataset/compute_norm_stats.py --patch-size 128

Run (env with ee+geemap+geopandas+rasterio):
  python patch_dataset/collect_glacier.py            # all polygons
  python patch_dataset/collect_glacier.py --limit 5  # smoke test
All tunables live in patch_dataset/config.py.
"""

import os
import sys
import argparse
import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import config as C
from collect_region import fetch_tile        # robust, terrain-correct downloader

import ee
import geopandas as gpd
import rasterio
from rasterio.features import rasterize


def main():
    ap = argparse.ArgumentParser(description="Enlarged per-glacier collection.")
    ap.add_argument("--shapefile", default=str(C.SHAPEFILE_PATH))
    ap.add_argument("--limit", type=int, default=None)
    args = ap.parse_args()

    try:
        ee.Initialize(project=C.GEE_PROJECT)
    except Exception:
        ee.Authenticate(); ee.Initialize(project=C.GEE_PROJECT)

    C.RAW_GLACIER_FEATURES_DIR.mkdir(parents=True, exist_ok=True)
    C.RAW_GLACIER_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"{len(gdf)} polygons | processing {n} | buffer={C.GLACIER_BUFFER_M:.0f} m | scale={C.SCALE_M:.0f} m")

    done = skip = 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"); skip += 1; continue

        fpath = C.RAW_GLACIER_FEATURES_DIR / f"feature_{idx:04d}.tif"
        mpath = C.RAW_GLACIER_MASKS_DIR / f"mask_{idx:04d}.tif"
        if fpath.exists() and mpath.exists():
            skip += 1; continue

        try:
            b = geom.bounds
            buf = C.GLACIER_BUFFER_M
            bbox = (b[0] - buf, b[1] - buf, b[2] + buf, b[3] + buf)
            arr, transform, crs = fetch_tile(bbox)
            H, W = arr.shape[1], arr.shape[2]
            with rasterio.open(fpath, "w", driver="GTiff", height=H, width=W,
                               count=arr.shape[0], dtype="float32",
                               crs=crs, transform=transform) as dst:
                dst.write(arr.astype(np.float32))
            mask = rasterize([(geom, 1)], out_shape=(H, W), transform=transform,
                             fill=0, dtype="uint8")
            with rasterio.open(mpath, "w", driver="GTiff", height=H, width=W, count=1,
                               dtype="uint8", crs=crs, transform=transform, nodata=0) as dst:
                dst.write(mask, 1)
            done += 1
            if done % 25 == 0 or idx == n - 1:
                print(f"  [{idx:04d}] done={done} skip={skip} fail={fail} (last {H}x{W}, fg={int(mask.sum())})")
        except Exception as e:
            fail += 1
            print(f"  [{idx:04d}] FAILED: {e}")

    print(f"\nPer-glacier collection complete. done={done} skipped={skip} failed={fail}")
    print(f"Features -> {C.RAW_GLACIER_FEATURES_DIR}")
    print(f"Masks    -> {C.RAW_GLACIER_MASKS_DIR}")


if __name__ == "__main__":
    main()