enrich_norway_static.py

#!/usr/bin/env python3
"""Fetch ship_type for Norway MMSIs missing from the snapshot LUT.

Uses VesselFinder's public popup-data endpoint (no auth, ~2 sec courtesy delay).
Endpoint: https://www.vesselfinder.com/api/pub/click/<MMSI>
Returns JSON like:
  {"type":"Oil Products Tanker","name":"HALTBAKK HULKEN","imo":9469209,
   "country":"Norway","gt":1543,"dw":2147,"al":78,"aw":12, ...}

Disk-cached at norway_ship_trajectory_datasets/data_raw/vesselfinder_cache.jsonl
so re-runs are free.
"""
from __future__ import annotations

import os
import argparse
import json
import sys
import time
from pathlib import Path
from urllib import request, error

ROOT = Path(os.environ.get("NORWAY_ROOT", "data/norway"))
CACHE = ROOT / "data_raw" / "vesselfinder_cache.jsonl"
LOG = ROOT / "logs" / f"vesselfinder_enrich_{time.strftime('%Y%m%dT%H%M%S')}.log"
URL_TMPL = "https://www.vesselfinder.com/api/pub/click/{mmsi}"
UA = "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 research-EnvShip-Bench"


def load_cache() -> dict[int, dict]:
    cache = {}
    if CACHE.exists():
        with CACHE.open("r") as fh:
            for ln in fh:
                ln = ln.strip()
                if not ln: continue
                try:
                    rec = json.loads(ln)
                    cache[int(rec["mmsi"])] = rec
                except Exception:
                    continue
    return cache


def append_cache(rec: dict) -> None:
    CACHE.parent.mkdir(parents=True, exist_ok=True)
    with CACHE.open("a") as fh:
        fh.write(json.dumps(rec, ensure_ascii=False) + "\n")


def fetch_vesselfinder(mmsi: int, timeout: int = 20) -> dict | None:
    req = request.Request(URL_TMPL.format(mmsi=mmsi), headers={"User-Agent": UA})
    try:
        with request.urlopen(req, timeout=timeout) as resp:
            return json.loads(resp.read().decode("utf-8"))
    except error.HTTPError as e:
        return {"_error": f"HTTP {e.code}"}
    except Exception as e:
        return {"_error": repr(e)[:200]}


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--mmsi-list", type=Path, required=True,
                   help="text file with one MMSI per line (the unknown set)")
    p.add_argument("--limit", type=int, default=0, help="cap N MMSIs (0=all)")
    p.add_argument("--delay", type=float, default=2.0, help="seconds between requests")
    p.add_argument("--retries-on-empty", type=int, default=0,
                   help="re-fetch MMSIs already cached as empty/error (default: skip)")
    args = p.parse_args()

    mmsis = [int(x.strip()) for x in args.mmsi_list.read_text().split() if x.strip().isdigit()]
    print(f"[enrich] {len(mmsis)} input MMSIs", flush=True)

    cache = load_cache()
    print(f"[enrich] cache contains {len(cache)} entries", flush=True)

    todo = []
    for m in mmsis:
        if m in cache and not args.retries_on_empty:
            continue
        if m in cache and ("_error" not in cache[m]) and cache[m].get("type"):
            continue
        todo.append(m)
    if args.limit:
        todo = todo[:args.limit]
    print(f"[enrich] {len(todo)} MMSIs to fetch", flush=True)

    LOG.parent.mkdir(parents=True, exist_ok=True)
    log_fh = LOG.open("w")
    t0 = time.time()
    n_ok = n_err = n_empty = 0
    for i, mmsi in enumerate(todo, 1):
        rec = fetch_vesselfinder(mmsi)
        if rec is None:
            rec = {"_error": "None"}
        rec["mmsi"] = mmsi
        rec["_ts"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
        append_cache(rec)
        if "_error" in rec:
            n_err += 1
        elif not rec.get("type"):
            n_empty += 1
        else:
            n_ok += 1
        if i % 50 == 0 or i == len(todo):
            elapsed = time.time() - t0
            rate = i / elapsed if elapsed else 0
            eta = (len(todo) - i) / rate if rate else float("inf")
            msg = f"[enrich] {i}/{len(todo)}  ok={n_ok} empty={n_empty} err={n_err}  rate={rate:.2f}/s  eta={eta/60:.1f}min"
            print(msg, flush=True)
            log_fh.write(msg + "\n"); log_fh.flush()
        time.sleep(args.delay)
    log_fh.close()
    print(f"[enrich] DONE  ok={n_ok}  empty={n_empty}  err={n_err}  cache={CACHE}", flush=True)


if __name__ == "__main__":
    main()