"""
RockGlacierNet — Region collection (whole-Sikkim grid).   [lives in patch_dataset/]

Grids the AOI (bounding box of all inventory polygons + AOI_MARGIN_M ~= the Sikkim
glacier belt) into RAW_TILE_PX tiles and downloads, per tile:
  * a 12-band Sentinel-2 + terrain GeoTIFF, and
  * a 1-band mask rasterized from ALL polygons onto that tile's exact grid.

REGION_MODE = "sampled" : every tile touching a glacier + REGION_BG_KEEP background tiles.
REGION_MODE = "dense"   : every tile in the AOI (full mosaic, multi-GB).

Robustness: the median composite is kept light with a scene-level cloud filter
(S2_CLOUD_PCT); if a tile still trips GEE "User memory limit exceeded", it is
auto-split into a finer sub-grid, each part downloaded and mosaicked in memory.

Outputs (nothing outside patch_dataset/):
  patch_dataset/raw_region/features/tile_RRRR_CCCC.tif
  patch_dataset/raw_region/masks/tile_RRRR_CCCC.tif

Run (env with ee+geemap+geopandas+rasterio, e.g. `brats`):
    python patch_dataset/collect_region.py                 # sampled (default)
    python patch_dataset/collect_region.py --mode dense
    python patch_dataset/collect_region.py --limit 3       # smoke test
All tunables live in patch_dataset/config.py.
"""

import os
import sys
import time
import argparse
import requests
import numpy as np

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

import ee
import geopandas as gpd
import rasterio
from rasterio.io import MemoryFile
from rasterio.merge import merge as rio_merge
from rasterio.features import rasterize
from shapely.geometry import box


# ─────────────────────────────────────────────────────────────────────────────
# Earth Engine: memory-light 12-band composite (band order == C.BAND_NAMES)
# ─────────────────────────────────────────────────────────────────────────────

def feature_image(roi):
    coll = (ee.ImageCollection(C.S2_COLLECTION)
            .filterBounds(roi)
            .filterDate(C.S2_DATE_START, C.S2_DATE_END)
            .filter(ee.Filter.lt("CLOUDY_PIXEL_PERCENTAGE", C.S2_CLOUD_PCT)))

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

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

    dem = ee.ImageCollection(C.DEM_COLLECTION).select("DEM").mosaic()
    # CRITICAL: reproject to a METRIC grid before ee.Terrain. After .mosaic() the
    # image defaults to EPSG:4326 (degrees); computing slope there yields garbage
    # (~1 deg everywhere). Reproject to UTM @30 m so slope/aspect are correct.
    dem_m = dem.reproject(crs=C.TARGET_CRS, scale=30)
    terr = ee.Terrain.products(dem_m)
    img = (s2.addBands(dem_m.rename("Elev"))
             .addBands(terr.select("slope").rename("Slope"))
             .addBands(terr.select("aspect").rename("Aspect")))
    return img.clip(roi)


# ─────────────────────────────────────────────────────────────────────────────
# Download helpers (UTM bbox -> GeoTIFF bytes, with memory-aware sub-tiling)
# ─────────────────────────────────────────────────────────────────────────────

def _utm_to_ee_roi(bbox_utm):
    ll = gpd.GeoSeries([box(*bbox_utm)], crs=C.TARGET_CRS).to_crs("EPSG:4326").total_bounds
    return ee.Geometry.Rectangle([ll[0], ll[1], ll[2], ll[3]])


def _download_bytes(bbox_utm, tries=3):
    """Download one bbox as GeoTIFF bytes. Raises MemoryError on GEE memory limit."""
    roi = _utm_to_ee_roi(bbox_utm)
    url = feature_image(roi).getDownloadURL(
        {"region": roi, "scale": C.SCALE_M, "crs": C.TARGET_CRS, "format": "GEO_TIFF"})
    last = None
    for k in range(tries):
        r = requests.get(url, timeout=300)
        if r.status_code == 200 and r.content[:2] in (b"II", b"MM"):   # valid TIFF magic
            return r.content
        last = r.content[:200]
        if r.status_code == 400 and b"memory limit" in r.content.lower():
            raise MemoryError(last)            # signal: split further
        time.sleep(2 * (k + 1))                # transient: retry
    raise RuntimeError(f"download failed: {last}")


def fetch_tile(bbox_utm, max_split=4):
    """Download bbox; on GEE memory error, recursively split into a finer grid and
    mosaic the parts in memory. Returns (array[C,H,W], transform, crs).

    NB: MemoryFile must stay open while its dataset is read, so we copy the array
    out before closing (otherwise GDAL reads a freed buffer)."""
    split = 1
    while True:
        minx, miny, maxx, maxy = bbox_utm
        xs = np.linspace(minx, maxx, split + 1)
        ys = np.linspace(miny, maxy, split + 1)
        try:
            chunks = [_download_bytes((xs[i], ys[j], xs[i+1], ys[j+1]))
                      for i in range(split) for j in range(split)]
        except MemoryError:
            if split >= max_split:
                raise RuntimeError("memory limit even at max split")
            split *= 2   # 1 -> 2 -> 4 (up to 16 sub-tiles)
            continue

        mfs, dsets = [], []
        try:
            for b in chunks:
                mf = MemoryFile(b); mfs.append(mf); dsets.append(mf.open())
            if len(dsets) == 1:
                ds = dsets[0]
                return ds.read().copy(), ds.transform, ds.crs
            arr, transform = rio_merge(dsets)
            return arr.copy(), transform, dsets[0].crs
        finally:
            for d in dsets: d.close()
            for mf in mfs: mf.close()


# ─────────────────────────────────────────────────────────────────────────────
# Main
# ─────────────────────────────────────────────────────────────────────────────

def main():
    ap = argparse.ArgumentParser(description="Whole-Sikkim region grid collection.")
    ap.add_argument("--shapefile", default=str(C.SHAPEFILE_PATH))
    ap.add_argument("--mode", choices=["sampled", "dense"], default=C.REGION_MODE)
    ap.add_argument("--limit", type=int, default=None, help="download at most N tiles")
    args = ap.parse_args()

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

    C.RAW_REGION_FEATURES_DIR.mkdir(parents=True, exist_ok=True)
    C.RAW_REGION_MASKS_DIR.mkdir(parents=True, exist_ok=True)

    gdf = gpd.read_file(args.shapefile).to_crs(C.TARGET_CRS)
    gdf = gdf[gdf.geometry.notna() & ~gdf.geometry.is_empty].copy()
    sidx = gdf.sindex

    minx, miny, maxx, maxy = gdf.total_bounds
    m, ts = C.AOI_MARGIN_M, C.RAW_TILE_PX * C.SCALE_M
    minx, miny, maxx, maxy = minx - m, miny - m, maxx + m, maxy + m
    xs = np.arange(minx, maxx, ts)
    ys = np.arange(miny, maxy, ts)
    rng = np.random.RandomState(C.RANDOM_SEED)

    # Plan tiles
    plan = []   # (row, col, bbox, has_glacier)
    for r, y in enumerate(ys):
        for c, x in enumerate(xs):
            bb = (x, y, x + ts, y + ts)
            cand = list(sidx.intersection(bb))
            has = bool(cand) and gdf.geometry.iloc[cand].intersects(box(*bb)).any()
            if args.mode == "dense" or has or rng.rand() < C.REGION_BG_KEEP:
                plan.append((r, c, bb, has))
    if args.limit:
        plan = plan[:args.limit]

    n_g = sum(1 for *_, h in plan if h)
    print(f"AOI {(maxx-minx)/1000:.0f}x{(maxy-miny)/1000:.0f} km | grid {len(xs)}x{len(ys)} "
          f"= {len(xs)*len(ys)} cells | mode={args.mode} | downloading {len(plan)} tiles "
          f"({n_g} glacier, {len(plan)-n_g} background)")

    done = skip = fail = 0
    for n, (r, c, bb, has) in enumerate(plan):
        fpath = C.RAW_REGION_FEATURES_DIR / f"tile_{r:04d}_{c:04d}.tif"
        mpath = C.RAW_REGION_MASKS_DIR / f"tile_{r:04d}_{c:04d}.tif"
        if fpath.exists() and mpath.exists():
            skip += 1; continue
        try:
            arr, transform, crs = fetch_tile(bb)
            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 all polygons intersecting this tile onto its grid
            cand = list(sidx.intersection(bb))
            geoms = [g for g in gdf.geometry.iloc[cand] if g.intersects(box(*bb))]
            mask = (rasterize([(g, 1) for g in geoms], out_shape=(H, W),
                              transform=transform, fill=0, dtype="uint8")
                    if geoms else np.zeros((H, W), "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 % 20 == 0 or n == len(plan) - 1:
                print(f"  {n+1}/{len(plan)}  done={done} skip={skip} fail={fail} "
                      f"(last tile {H}x{W}, fg={int(mask.sum())})")
        except Exception as e:
            fail += 1
            print(f"  tile_{r:04d}_{c:04d} FAILED: {e}")

    print(f"\nRegion collection complete. done={done} skipped={skip} failed={fail}")
    print(f"Features -> {C.RAW_REGION_FEATURES_DIR}")
    print(f"Masks    -> {C.RAW_REGION_MASKS_DIR}")


if __name__ == "__main__":
    main()