hf_data/basic.py

"""HF-backed basic-data adapter for the AQUAH->CREST_demo integration.

Replaces AQUAH's live basic-data download + clip: windowed-reads the HydroSHEDS
DEM / flow-dir / flow-acc COGs hosted in `vincewin/CREST_data` and writes the
three EF5 inputs (dem_clip.tif, fdir_clip.tif, facc_clip.tif) for a basin bbox.

All three HydroSHEDS grids share one grid, so clipping the same bbox yields
pixel-aligned outputs (EF5 masks the catchment internally from the outlet cell).

On this network GDAL's curl needs GDAL_HTTP_UNSAFESSL=YES (corporate MITM);
that is LOCAL-ONLY and not set on HF cloud runners.
"""
from __future__ import annotations

import os
import time
from dataclasses import dataclass

import rasterio
from rasterio.windows import from_bounds

HF_REPO = "vincewin/CREST_data"
_RESOLVE = f"https://huggingface.co/datasets/{HF_REPO}/resolve/main"

# EF5 input filename <- COG path in CREST_data
BASIC_GRIDS = {
    "dem_clip.tif": "basic/na_con_3s.tif",
    "fdir_clip.tif": "basic/na_dir_3s.tif",
    "facc_clip.tif": "basic/na_acc_3s.tif",
}


def _vsicurl(cog_path: str) -> str:
    return f"/vsicurl/{_RESOLVE}/{cog_path}"


def _cog_src(cog: str, force_download: bool = False) -> str:
    """Local whole-file copy when cached (or force_download) else /vsicurl.

    Windowed /vsicurl reads are anonymous range requests — under fleet-scale
    concurrency HF rate-limits them (RasterioIOError storms across every
    runner), so a local copy (authenticated hf_hub download, shared across
    processes via HF_HOME) is preferred whenever it exists."""
    try:
        from huggingface_hub import hf_hub_download, try_to_load_from_cache
        p = try_to_load_from_cache(HF_REPO, cog, repo_type="dataset")
        if isinstance(p, str) and os.path.exists(p):
            return p
        if force_download:
            return hf_hub_download(HF_REPO, cog, repo_type="dataset",
                                   token=os.environ.get("HF_TOKEN"))
    except Exception:
        pass
    return _vsicurl(cog)


def prefetch() -> None:
    """Best-effort download of the three CONUS COGs (~6.9 GB) into the HF
    cache so every later clip reads locally. Bulk runners (the fleet) call
    this once at startup; interactive use never needs it."""
    from huggingface_hub import hf_hub_download
    for cog in BASIC_GRIDS.values():
        try:
            hf_hub_download(HF_REPO, cog, repo_type="dataset",
                            token=os.environ.get("HF_TOKEN"))
        except Exception:
            pass


@dataclass
class ClipResult:
    out_dir: str
    files: dict[str, str]
    width: int
    height: int
    bounds: tuple[float, float, float, float]
    derived: bool = False        # dir/acc regenerated from the DEM (pysheds)


# ESRI D8 codes EF5 expects with ESRIDDM=true (E,SE,S,SW,W,NW,N,NE)
_ESRI_D8 = (1, 2, 4, 8, 16, 32, 64, 128)


def _dir_acc_usable(fdir_path: str, facc_path: str, dem_path: str,
                    max_bad_frac: float = 0.05) -> bool:
    """Sanity-check the clipped flow-dir/acc: on cells where the DEM is valid,
    the direction must be an ESRI D8 code and the accumulation non-negative."""
    import numpy as np
    with rasterio.open(dem_path) as d, rasterio.open(fdir_path) as f, \
            rasterio.open(facc_path) as a:
        dem, fdir, facc = d.read(1), f.read(1), a.read(1)
    ok_dem = dem != -9999.0
    if not ok_dem.any():
        return False
    good_dir = np.isin(fdir, _ESRI_D8)
    good_acc = facc >= 0
    bad = ok_dem & ~(good_dir & good_acc)
    return (bad.sum() / ok_dem.sum()) <= max_bad_frac


def derive_dir_acc(dem_path: str, fdir_path: str, facc_path: str) -> None:
    """Fallback: derive D8 flow direction + accumulation from the clipped DEM
    with pysheds (condition -> flowdir -> accumulation). pysheds' default
    dirmap (N,NE,E,SE,S,SW,W,NW)=(64,128,1,2,4,8,16,32) IS the ESRI encoding."""
    import numpy as np
    from pysheds.grid import Grid

    grid = Grid.from_raster(dem_path)
    dem = grid.read_raster(dem_path)
    dem_f = grid.fill_pits(dem)
    dem_f = grid.fill_depressions(dem_f)
    dem_f = grid.resolve_flats(dem_f)
    fdir = grid.flowdir(dem_f)
    facc = grid.accumulation(fdir)

    with rasterio.open(dem_path) as ref:
        profile = ref.profile.copy()
    profile.update(dtype="float32", nodata=-9999.0,
                   tiled=False, blockysize=1, crs="EPSG:4326",
                   interleave="pixel")
    profile.pop("compress", None)
    mask = np.asarray(dem) == -9999.0
    fdir = np.asarray(fdir, dtype="float32")
    fdir[~np.isin(fdir, _ESRI_D8)] = -9999.0   # pysheds pits/flats: -1/-2/0
    facc = np.asarray(facc, dtype="float32")
    for path, out in ((fdir_path, fdir), (facc_path, facc)):
        out[mask] = -9999.0
        with rasterio.open(path, "w", **profile) as dst:
            dst.write(out, 1)


def store_dir(bbox) -> str:
    """Shared per-basin clip store: repeated runs (and every calibration
    candidate) reuse the same clipped DEM/DDM/FAM instead of re-reading the
    remote COGs — faster AND immune to transient network read failures."""
    from hf_data.statecache import CACHE_DIR
    key = "_".join(f"{v:.3f}" for v in bbox).replace("-", "m").replace(".", "p")
    d = os.path.join(CACHE_DIR, "basic", key)
    os.makedirs(d, exist_ok=True)
    return d


def clip_basic_data(bbox, out_dir: str, unsafe_ssl: bool | None = None) -> ClipResult:
    """Clip DEM/DDM/FAM to `bbox` (W, S, E, N in EPSG:4326) from the HF COGs.

    Writes dem_clip.tif / fdir_clip.tif / facc_clip.tif into `out_dir`, all on
    the same window/geotransform. Returns a ClipResult. If all three clips are
    already present in `out_dir` (shared store), they are reused as-is.
    """
    W, S, E, N = bbox
    # local-network MITM workaround; harmless (and unset) on HF runners
    if unsafe_ssl is None:
        unsafe_ssl = os.name == "nt"
    if unsafe_ssl:
        os.environ.setdefault("GDAL_HTTP_UNSAFESSL", "YES")

    os.makedirs(out_dir, exist_ok=True)
    files: dict[str, str] = {}
    ref_shape = None
    ref_bounds = None

    existing = {n: os.path.join(out_dir, n) for n in BASIC_GRIDS}
    if all(os.path.exists(p) for p in existing.values()):     # store hit
        with rasterio.open(existing["dem_clip.tif"]) as ds:
            t = ds.transform
            return ClipResult(out_dir=out_dir, files=existing,
                              width=ds.width, height=ds.height,
                              bounds=(t.c, t.f + ds.height * t.e,
                                      t.c + ds.width * t.a, t.f),
                              derived=False)

    fetch_failed: list[str] = []
    for out_name, cog in BASIC_GRIDS.items():
        err = None
        for attempt in range(4):               # /vsicurl reads fail transiently;
            try:                               # final tries clip a local copy
                _clip_one(_cog_src(cog, force_download=attempt >= 2),
                          out_name, bbox, out_dir, files)
                err = None
                break
            except Exception as e:
                err = e
                time.sleep(2 * (attempt + 1))
        if err is not None:
            if out_name == "dem_clip.tif":
                raise err                      # DEM is required — no fallback
            fetch_failed.append(out_name)
    _shapes = {}
    for out_name, path in files.items():
        with rasterio.open(path) as ds:
            _shapes[out_name] = (ds.height, ds.width)
            if out_name == "dem_clip.tif":
                t = ds.transform
                ref_shape = (ds.height, ds.width)
                ref_bounds = (t.c, t.f + ds.height * t.e,
                              t.c + ds.width * t.a, t.f)
    for out_name, shp in _shapes.items():
        if shp != ref_shape:
            fetch_failed.append(out_name)      # misaligned -> rederive

    # fallback: if dir/acc failed to fetch, are misaligned, or look broken
    # (bad D8 codes, negative accumulation), regenerate from the clipped DEM
    derived = False
    fdir_p = os.path.join(out_dir, "fdir_clip.tif")
    facc_p = os.path.join(out_dir, "facc_clip.tif")
    try:
        usable = (not fetch_failed) and _dir_acc_usable(
            files["fdir_clip.tif"], files["facc_clip.tif"], files["dem_clip.tif"])
    except Exception:
        usable = False
    if not usable:
        derive_dir_acc(files["dem_clip.tif"], fdir_p, facc_p)
        files["fdir_clip.tif"], files["facc_clip.tif"] = fdir_p, facc_p
        derived = True

    return ClipResult(out_dir=out_dir, files=files,
                      width=ref_shape[1], height=ref_shape[0], bounds=ref_bounds,
                      derived=derived)


def _clip_one(src_path: str, out_name: str, bbox, out_dir: str, files: dict) -> None:
    W, S, E, N = bbox
    with rasterio.open(src_path) as src:
        win = from_bounds(W, S, E, N, src.transform).round_offsets().round_lengths()
        # standardized clip for EF5's TifGrid reader: single-band Float32,
        # strip-organized, nodata=-9999, WGS84
        data = src.read(1, window=win).astype("float32")
        if src.nodata is not None:
            data[data == float(src.nodata)] = -9999.0
        transform = src.window_transform(win)
        profile = src.profile.copy()
        profile.update(
            driver="GTiff", height=data.shape[0], width=data.shape[1],
            transform=transform, tiled=False, blockysize=1,
            dtype="float32", nodata=-9999.0, crs="EPSG:4326",
            interleave="pixel",         # contiguous — the COGs are band-
                                        # interleaved, which EF5's reader
                                        # (TIFFReadScanline sample arg) rejects
        )
        profile.pop("compress", None)   # uncompressed 1-row strips, like
                                        # EF5's own WriteFloatTifGrid output
    out_path = os.path.join(out_dir, out_name)
    with rasterio.open(out_path, "w", **profile) as dst:
        dst.write(data, 1)
    files[out_name] = out_path


if __name__ == "__main__":
    import sys, time
    # smoke test: Allagash gauge 01011000 basin bbox
    bbox = (-69.83, 46.32, -68.33, 47.82)
    out = sys.argv[1] if len(sys.argv) > 1 else "._clip_test"
    t = time.time()
    r = clip_basic_data(bbox, out)
    print(f"clipped {r.width}x{r.height} in {time.time()-t:.1f}s -> {out}")
    for name, path in r.files.items():
        with rasterio.open(path) as ds:
            a = ds.read(1)
            nod = ds.nodata
            valid = a[a != nod] if nod is not None else a
            vmin = float(valid.min()) if valid.size else float("nan")
            vmax = float(valid.max()) if valid.size else float("nan")
            print(f"  {name}: {ds.dtypes[0]} nodata={nod} range[{vmin:.1f},{vmax:.1f}]")