Northern-Thailand-Ping-River-Monitor/src/hii_backfill.py

"""Backfill historical water levels from the HII waterlevel_graph endpoint.

The api-v3 waterlevel_graph archive reaches back to ~2019 with hourly
wl_msl + discharge. This module walks a date range in chunks per station and
upserts into hii_waterlevel (idempotent; safe to re-run and to overlap with
the live snapshot collector). Station metadata comes from a live
waterlevel_load fetch, so hii_wl_stations is populated/refreshed as a side
effect.

Usage: python scripts/backfill_hii_waterlevel.py --start 2019-01-01
"""

import argparse
import datetime
import logging
import time
from typing import Dict, List, Optional

from .hii_collector import (
    PING_BASIN_CODE,
    HiiClient,
    HiiStore,
    _parse_datetime,
    _to_float,
)

logger = logging.getLogger(__name__)

DEFAULT_START = datetime.date(2019, 1, 1)
DEFAULT_CHUNK_DAYS = 365  # full-year windows verified working (8,760 rows, ~700KB)
DEFAULT_SLEEP_SECONDS = 1.0


def parse_graph_rows(payload: Dict) -> List[Dict]:
    """Extract history rows from a waterlevel_graph payload (skips empty rows)."""
    rows = (payload.get("data") or {}).get("graph_data") or []
    records = []
    for row in rows:
        timestamp = _parse_datetime(row.get("datetime"))
        wl_msl = _to_float(row.get("value"))
        discharge = _to_float(row.get("discharge"))
        if timestamp is None or (wl_msl is None and discharge is None):
            continue
        records.append(
            {"timestamp": timestamp, "wl_msl": wl_msl, "discharge": discharge}
        )
    return records


def fetch_waterlevel_history(
    client: HiiClient,
    station_id: int,
    start_date: datetime.date,
    end_date: datetime.date,
) -> List[Dict]:
    """Hourly wl_msl + discharge history (archive reaches back to ~2019)."""
    payload = client.get(
        "waterlevel_graph",
        params={
            "station_type": "tele_waterlevel",
            "station_id": station_id,
            "start_date": start_date.isoformat(),
            "end_date": end_date.isoformat(),
        },
    )
    return parse_graph_rows(payload)


def chunk_date_range(
    start: datetime.date, end: datetime.date, chunk_days: int
) -> List[tuple]:
    """Split [start, end] into inclusive (start, end) windows."""
    chunks = []
    cursor = start
    while cursor <= end:
        chunk_end = min(cursor + datetime.timedelta(days=chunk_days - 1), end)
        chunks.append((cursor, chunk_end))
        cursor = chunk_end + datetime.timedelta(days=1)
    return chunks


def select_stations(
    station_records: List[Dict],
    codes: Optional[List[str]] = None,
    all_stations: bool = False,
) -> List[Dict]:
    """Pick stations to backfill from parsed waterlevel_load records.

    Default: stations that mirror a RID gauge (rid_code) or are flagged
    is_key_station — the ones relevant to the flood model. Explicit codes
    match rid_code or oldcode; --all takes every station in the basin.
    """
    if all_stations:
        return station_records
    if codes:
        wanted = {c.strip().upper() for c in codes if c.strip()}
        return [
            r
            for r in station_records
            if (r.get("rid_code") or "").upper() in wanted
            or (r.get("oldcode") or "").upper() in wanted
        ]
    return [r for r in station_records if r.get("rid_code") or r.get("is_key_station")]


def backfill(
    store: HiiStore,
    client: Optional[HiiClient] = None,
    start: datetime.date = DEFAULT_START,
    end: Optional[datetime.date] = None,
    codes: Optional[List[str]] = None,
    all_stations: bool = False,
    chunk_days: int = DEFAULT_CHUNK_DAYS,
    sleep_seconds: float = DEFAULT_SLEEP_SECONDS,
    basin_code: int = PING_BASIN_CODE,
) -> Dict[str, int]:
    """Run the backfill; returns {'stations': n, 'rows': n, 'errors': n}."""
    client = client or HiiClient()
    end = end or datetime.date.today()

    logger.info("Fetching station catalog from waterlevel_load...")
    station_records = client.fetch_waterlevel(basin_code)
    # Refresh station metadata (and today's snapshot) while we have it
    store.save_waterlevel(station_records)

    stations = select_stations(station_records, codes=codes, all_stations=all_stations)
    if not stations:
        logger.error("No stations matched the selection")
        return {"stations": 0, "rows": 0, "errors": 0}

    chunks = chunk_date_range(start, end, chunk_days)
    logger.info(
        f"Backfilling {len(stations)} stations x {len(chunks)} windows "
        f"({start} .. {end}, {chunk_days}-day chunks)"
    )

    totals = {"stations": len(stations), "rows": 0, "errors": 0}
    for station in stations:
        sid = station["station_id"]
        label = station.get("rid_code") or station.get("oldcode") or str(sid)
        station_rows = 0
        for chunk_start, chunk_end in chunks:
            try:
                rows = fetch_waterlevel_history(client, sid, chunk_start, chunk_end)
                station_rows += store.save_waterlevel_history(sid, rows)
            except Exception as e:
                totals["errors"] += 1
                logger.warning(f"{label}: {chunk_start}..{chunk_end} failed: {e}")
            time.sleep(sleep_seconds)
        totals["rows"] += station_rows
        logger.info(f"{label} (id {sid}): {station_rows} rows saved")

    logger.info(
        f"Backfill complete: {totals['rows']} rows across "
        f"{totals['stations']} stations, {totals['errors']} failed windows"
    )
    return totals


def main(argv: Optional[List[str]] = None) -> bool:
    parser = argparse.ArgumentParser(
        description="Backfill hii_waterlevel from the HII waterlevel_graph archive"
    )
    parser.add_argument(
        "--start",
        type=datetime.date.fromisoformat,
        default=DEFAULT_START,
        help=f"First date to fetch (default {DEFAULT_START})",
    )
    parser.add_argument(
        "--end",
        type=datetime.date.fromisoformat,
        default=None,
        help="Last date to fetch (default today)",
    )
    parser.add_argument(
        "--stations",
        help="Comma-separated codes (rid_code or oldcode, e.g. P.1,P.67,CHM004). "
        "Default: all RID-mirror and key stations",
    )
    parser.add_argument(
        "--all",
        action="store_true",
        help="Backfill every Ping-basin station (125+; slow)",
    )
    parser.add_argument(
        "--chunk-days", type=int, default=DEFAULT_CHUNK_DAYS, help="Window size"
    )
    parser.add_argument(
        "--sleep",
        type=float,
        default=DEFAULT_SLEEP_SECONDS,
        help="Pause between requests in seconds",
    )
    args = parser.parse_args(argv)

    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
    )

    from .config import database_url

    url = database_url()
    if not url:
        logger.error("no database: set POSTGRES_CONNECTION_STRING or POSTGRES_*")
        return False
    store = HiiStore(url)
    if not store.connect():
        return False

    totals = backfill(
        store,
        start=args.start,
        end=args.end,
        codes=args.stations.split(",") if args.stations else None,
        all_stations=args.all,
        chunk_days=args.chunk_days,
        sleep_seconds=args.sleep,
    )
    return totals["rows"] > 0 and totals["errors"] == 0