thermograph/backend/seed_era5.py

"""One-time ERA5 seed for the curated-city cells (keeps ERA5 fidelity without Open-Meteo).

Now that the live history primary is NASA POWER, this backfills the high-traffic
curated-city cells (backend/cities.json) with true ERA5 data pulled from a public,
KEYLESS source — the Earthmover Icechunk ERA5 archive on AWS Open Data — so the
pages people actually visit keep ERA5-quality history. It writes straight into the
history store (bypassing the live fetch), so a seeded cell is served from cache and
never falls to NASA on first request.

    # on a machine with network + the seed deps (see requirements-seed.txt):
    pip install -r requirements-seed.txt
    python seed_era5.py --dry-run            # fetch+transform one cell, print, no write
    python seed_era5.py [--limit N] [--overwrite]

Idempotent by default: a cell that already has a cached history is skipped (pass
--overwrite to reseed). Not run in CI or the app image — the icechunk/xarray/zarr
stack lives only in requirements-seed.txt.

IMPORTANT — the S3/Icechunk *access* layer (_open_store / fetch_point) targets a
young, fast-moving API and the store's exact variable/coordinate names; it is NOT
exercised by the test suite. VERIFY it with `--dry-run` on the seed box before a
full run. The hourly→daily transform (hourly_to_daily) IS unit-tested and is the
part that must be numerically correct.

Source: https://registry.opendata.aws/earthmover-era5/  (group="temporal", anonymous)
"""
import os
import sys

import polars as pl

from data import cities
from data import climate
from data import grid

# Match the live archive span (climate.START_DATE / ARCHIVE_LATENCY_DAYS).
START_DATE = climate.START_DATE            # "1980-01-01"

# Anonymous AWS Open Data Icechunk store; overridable in case the registry moves.
ERA5_BUCKET = os.environ.get("THERMOGRAPH_ERA5_BUCKET", "earthmover-icechunk-era5")
ERA5_PREFIX = os.environ.get("THERMOGRAPH_ERA5_PREFIX", "")   # verify on the seed box
ERA5_REGION = os.environ.get("THERMOGRAPH_ERA5_REGION", "us-east-1")
ERA5_GROUP = os.environ.get("THERMOGRAPH_ERA5_GROUP", "temporal")  # time-contiguous layout

# ERA5 store variable name -> our short name. ERA5 gusts (i10fg) are real here, so the
# seed keeps genuine measured gusts rather than the Meteostat/estimate fallback.
ERA5_VARS = {
    "2m_temperature": "t2m",                 # K
    "2m_dewpoint_temperature": "d2m",        # K (→ RH with t2m)
    "total_precipitation": "tp",             # m, hourly accumulation
    "10m_u_component_of_wind": "u10",        # m/s
    "10m_v_component_of_wind": "v10",        # m/s
    "instantaneous_10m_wind_gust": "i10fg",  # m/s
}

_M_TO_IN = 39.37007874
_MS_TO_MPH = 2.2369362920544


def hourly_to_daily(hourly: pl.DataFrame) -> pl.DataFrame:
    """Aggregate an hourly ERA5 point frame to the daily schema climate.py stores,
    in the app's units (°F, inches, mph, %RH). Input columns: time (Datetime, UTC),
    t2m/d2m (K), tp (m, per-hour accumulation), u10/v10/i10fg (m/s).

    Days are UTC calendar days — the reanalysis is UTC, so this is a small offset
    from a local-timezone daily boundary, immaterial for percentile climatology.
    Relative humidity is the Magnus-formula RH from temperature and dewpoint."""
    tc = pl.col("t2m") - 273.15
    tdc = pl.col("d2m") - 273.15
    es = (17.625 * tc / (243.04 + tc)).exp()
    e = (17.625 * tdc / (243.04 + tdc)).exp()
    df = hourly.with_columns(
        pl.col("time").dt.date().alias("date"),
        ((pl.col("t2m") - 273.15) * 9.0 / 5.0 + 32.0).alias("t_f"),
        ((pl.col("u10") ** 2 + pl.col("v10") ** 2).sqrt() * _MS_TO_MPH).alias("wind_mph"),
        (pl.col("i10fg") * _MS_TO_MPH).alias("gust_mph"),
        (pl.col("tp") * _M_TO_IN).alias("precip_in"),
        (100.0 * e / es).clip(0.0, 100.0).alias("rh"),
    )
    return (
        df.group_by("date")
        .agg(
            pl.col("t_f").max().alias("tmax"),
            pl.col("t_f").min().alias("tmin"),
            pl.col("precip_in").sum().alias("precip"),
            pl.col("wind_mph").max().alias("wind"),
            pl.col("gust_mph").max().alias("gust"),
            pl.col("rh").mean().alias("humid"),
        )
        .sort("date")
    )


def _default_end() -> str:
    import datetime
    return (datetime.date.today()
            - datetime.timedelta(days=climate.ARCHIVE_LATENCY_DAYS)).isoformat()


def _open_store():
    """Open the anonymous Icechunk ERA5 store (temporal group). Lazy-imports the
    seed-only deps. VERIFY against the current icechunk/xarray API on the seed box —
    this young API and the store's exact paths are not covered by tests."""
    import icechunk           # noqa: PLC0415 - seed-only dep (requirements-seed.txt)
    import xarray as xr       # noqa: PLC0415

    storage = icechunk.s3_storage(
        bucket=ERA5_BUCKET, prefix=ERA5_PREFIX, region=ERA5_REGION, anonymous=True)
    repo = icechunk.Repository.open(storage)
    session = repo.readonly_session("main")
    return xr.open_zarr(session.store, group=ERA5_GROUP, consolidated=False)


def fetch_point(ds, lat: float, lon: float, start: str, end: str) -> pl.DataFrame:
    """Pull the hourly ERA5 series for the nearest grid point over [start, end] into a
    polars frame (no pandas). ERA5 longitudes are 0..360. Coordinate/variable names
    may differ per store revision — verify with --dry-run."""
    sub = (ds[list(ERA5_VARS)]
           .sel(latitude=lat, longitude=lon % 360.0, method="nearest")
           .sel(time=slice(start, end))
           .load())
    data = {"time": sub["time"].values}
    for src, dst in ERA5_VARS.items():
        data[dst] = sub[src].values
    return pl.DataFrame(data)


def seed_cell(ds, cell: dict, start: str, end: str) -> int:
    """Fetch → daily → finalize → write one cell's ERA5 history. Returns row count."""
    daily = hourly_to_daily(fetch_point(ds, cell["center_lat"], cell["center_lon"], start, end))
    # ERA5 has no apparent-temperature field; approximate fmax/fmin from the NWS heat
    # index / wind chill (same as the NASA/MET paths). _finalize_approximated adds
    # feels + doy and folds NaN→null; the write path strips doy.
    frame = climate._finalize_approximated(daily)
    climate._write_history_backed(cell["id"], frame)
    return frame.height


def main(limit: "int | None" = None, overwrite: bool = False,
         dry_run: bool = False, start: str = START_DATE, end: "str | None" = None) -> None:
    end = end or _default_end()
    ds = _open_store()
    todo = cities.all_cities()
    if limit:
        todo = todo[:limit]
    seeded = skipped = failed = 0
    for i, c in enumerate(todo, 1):
        cell = grid.snap(c["lat"], c["lon"])
        if not overwrite:
            cached = climate.load_cached_history(cell)
            if cached is not None and not cached.is_empty():
                skipped += 1
                continue
        try:
            if dry_run:
                daily = hourly_to_daily(
                    fetch_point(ds, cell["center_lat"], cell["center_lon"], start, end))
                print(daily.head())
                print(f"[dry-run] {c['slug']} ({cell['id']}): {daily.height} days "
                      f"({start}..{end}); nothing written")
                return
            n = seed_cell(ds, cell, start, end)
            seeded += 1
            print(f"[{i}/{len(todo)}] seeded {c['slug']} ({cell['id']}): {n} days")
        except Exception as e:  # noqa: BLE001 - keep going; a failed cell self-heals via NASA
            failed += 1
            print(f"[{i}/{len(todo)}] FAILED {c['slug']}: {e}")
    print(f"done: seeded={seeded} skipped(cached)={skipped} failed={failed}")


if __name__ == "__main__":
    argv = sys.argv[1:]
    lim = int(argv[argv.index("--limit") + 1]) if "--limit" in argv else None
    main(limit=lim, overwrite="--overwrite" in argv, dry_run="--dry-run" in argv)