#!/usr/bin/env python3
"""
Attach the places each work is about to overview/data.json and write overview/places.json.

Usage:
    python build_overview.py && python geolocate.py

Where places come from
- Gutenberg: LoC subject headings. "X -- Description and travel" / "X -- Discovery and
  exploration" name X; geographic subdivisions such as "Birds -- Great Britain" name the
  place after the topic.
- All other sources, and Gutenberg works whose subject headings name no place: an LLM on
  innkube (qwen38-27b, Uni Passau) reads the title and, for Wikisource and Humboldt's
  journals, the short description (for BHL, the Internet Archive and Gallica, the subject
  headings), and names the places as modern English names.
  This handles adjectives and historical names ("persischen Hofe" -> Iran,
  "Morgenland" -> Orient). Answers are cached in llm_places.json. Needs INNKUBE_API_KEY.

Geocoding
- Broad or historical regions that OpenStreetMap has no single point for (Africa, Central;
  Arctic regions; Morgenland; Ostindien …) use the approximate centres in REGIONS.
- Everything else goes to Nominatim (OpenStreetMap), one request per second, cached in
  geocache.json so reruns only ask for new names. Only results that are places (countries,
  regions, cities, islands, rivers, mountains, seas) are kept.
"""
import argparse
import hashlib
import json
import os
import re
import sys
import time
import urllib.parse
import urllib.request
from collections import Counter
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

NOMINATIM = "https://nominatim.openstreetmap.org/search"
INNKUBE = "https://llms.innkube.fim.uni-passau.de/v1/chat/completions"
LLM_MODEL = "qwen38-27b"
LLM_PROMPT = (
    "For each numbered title (German or English) of a historical travel account or natural-history text, list the "
    "geographic places the text is ABOUT (not where it was printed). Use modern English names, most "
    "specific first, as 'Place, Country' where useful (e.g. 'Allgäu Alps, Germany'). For broad "
    "historical regions use one of: Orient, Middle East, East Indies, West Indies, Arctic regions, "
    "Africa, Central; Africa, West; Africa, East; South Sea Islands; Siberia. If no place is named or "
    "clearly implied, return an empty list. Answer only with JSON mapping each number to a list of "
    "strings.\n\n")
UA = {"User-Agent": "nature-travel-catalogue/1.0 (https://thomashaider.net/natural-histories/)"}

# Approximate centres for regions without a single OSM point. (lat, lon)
REGIONS = {
    # ancient and biblical regions, which Nominatim either lacks or matches to American towns
    "assyria": (36.4, 43.2), "mesopotamia": (33.0, 44.0), "canaan": (32.0, 35.2), "holy land": (31.8, 35.2),
    "palestine": (31.9, 35.2), "babylonia": (32.5, 44.4), "phoenicia": (33.9, 35.5), "judea": (31.7, 35.1),
    "jordan river": (32.3, 35.55), "river jordan": (32.3, 35.55), "jordan valley": (32.3, 35.55),
    "jordan river, jordan": (32.3, 35.55), "jordan river, israel": (32.3, 35.55),
    "africa": (2, 20), "africa, central": (2, 22), "central africa": (2, 22), "west africa": (10, -3), "east africa": (0, 37), "north africa": (28, 10), "africa, east": (0, 37), "africa, west": (10, -3),
    "africa, north": (28, 10), "africa, south": (-27, 25), "africa, southern": (-25, 25),
    "africa, northeast": (12, 40), "africa, southwest": (-22, 17), "america": (15, -85),
    "north america": (45, -100), "south america": (-15, -60), "central america": (13, -86),
    "latin america": (0, -70), "west (u.s.)": (40, -112), "southwest, new": (34, -108),
    "northwest, pacific": (46, -121), "northwest, canadian": (60, -115), "southern states": (33, -86),
    "new england": (43.5, -71.5), "middle west": (41, -90), "great plains": (40, -100),
    "rocky mountains": (43, -110), "middle east": (29, 42), "orient": (30, 45), "morgenland": (30, 40),
    "levant": (33, 36), "holy land": (31.5, 35), "east asia": (35, 115), "far east": (35, 120),
    "asia": (40, 90), "asia, central": (43, 65), "asia, southeastern": (8, 110), "asia minor": (39, 32),
    "oceania": (-10, 160), "south sea islands": (-15, -150), "islands of the pacific": (-10, 170),
    "pacific ocean": (0, -160), "atlantic ocean": (25, -40), "indian ocean": (-20, 75),
    "arctic regions": (75, -40), "arktis": (78, 0), "polar regions": (80, 0), "north pole": (88, 0),
    "nordpol": (88, 0), "antarctica": (-80, 0), "antarctic regions": (-75, 0), "südpol": (-88, 0),
    "europe": (50, 15), "europa": (50, 15), "europe, western": (48, 3), "europe, eastern": (50, 30),
    "europe, central": (49, 13), "europe, northern": (62, 15), "scandinavia": (63, 15),
    "skandinavien": (63, 15), "balkan peninsula": (42, 22), "mediterranean region": (37, 15),
    "mittelmeer": (37, 15), "alps": (46.5, 10), "alpen": (46.5, 10), "sahara": (23, 10),
    "siberia": (62, 100), "sibirien": (62, 100), "west indies": (18, -70), "westindien": (18, -70),
    "east indies": (0, 115), "ostindien": (15, 80), "ost-indien": (15, 80), "indien": (22, 79),
    "hindostan": (25, 80), "abessinien": (9, 39), "persien": (32, 53), "siam": (15, 101),
    "ceylon": (7.8, 80.7), "formosa": (23.7, 121), "lappland": (68, 22), "orient (morgenland)": (30, 40),
    "amerika": (15, -85), "nordamerika": (45, -100), "südamerika": (-15, -60), "mittelamerika": (13, -86),
    "afrika": (2, 20), "ostafrika": (0, 37), "westafrika": (10, -3), "südafrika": (-29, 25),
    "nordafrika": (28, 10), "zentralafrika": (2, 22), "innerafrika": (2, 22), "asien": (40, 90),
    "australien": (-25, 134), "ozeanien": (-10, 160), "südsee": (-15, -150), "tropics": (0, 0),
    "deutsch-ostafrika": (-6, 35), "deutsch-südwestafrika": (-22, 17), "kamerun": (5.7, 12.4),
    "levante": (33, 36), "kleinasien": (39, 32),
    # historical regions that Nominatim matches to same-named villages
    "nubia": (21, 32), "pampas": (-35, -63), "rhodesia": (-19, 30), "transvaal": (-25, 29),
    "provence": (43.9, 6), "riviera": (43.7, 7.3), "keewatin": (63, -94),
    "selborne": (51.098, -0.942),
    # LoC: plain "Georgia" is the US state; the country is "Georgia (Republic)"
    "georgia": (32.7, -83.4), "palästina": (31.5, 35), "holy land (palestine)": (31.5, 35),
    # LLM names that Nominatim matched to villages, towns or industrial sites of the same name
    "northwest, united states": (46, -121), "northwest, canada": (60, -115), "west (u.s.), united states": (40, -112),
    "west, united states": (40, -112), "hudson bay": (60, -85), "hudson bay, canada": (60, -85),
    "patagonia": (-45, -69), "patagonia, argentina": (-45, -69), "kashmir": (34.1, 74.8), "kashmir, india": (34.1, 74.8),
    "south america, south america": (-15, -60), "dunk island, australia": (-17.94, 146.14), "burma": (21, 96),
    "burma, myanmar": (21, 96), "myanmar, asia": (21, 96), "china, china": (35, 103), "tartary": (45, 80),
    "mexico, north america": (23, -102), "lake victoria, uganda": (-1, 33), "acadia, canada": (45.5, -64),
    "equatorial regions": (0, 20), "argentina, argentina": (-34, -64), "china seas": (15, 115),
    "puerto rico, puerto rico": (18.2, -66.5), "jersey, united kingdom": (49.21, -2.13), "nyasaland, malawi": (-13.5, 34),
    "livonia": (57.5, 25), "pacific coast, canada": (52, -128), "lochaber, scotland": (56.85, -5.1),
    "colombia, colombia": (4, -73), "red river settlement, canada": (49.9, -97.1), "indian archipelago": (0, 115),
    "new holland": (-25, 134), "persia": (32, 53), "afghanistan, asia": (34, 66), "philippines, asia": (12, 122),
    "saint helena, south atlantic": (-15.97, -5.71),
}
# LoC subdivisions and topical headings that are not places.
NOT_PLACE = re.compile(
    r"^(description|discovery|juvenile|early works|pictorial|history|periodicals|guidebooks|"
    r"biography|identification|classification|anecdotes|bibliograph|social life|maps|views|"
    r"fiction|poetry|congresses|catalogs|study and teaching|habits|geographical distribution|"
    r"collection|ecology|anatomy|behavior|folklore|dictionaries|handbooks|antiquities|"
    r"natural history|travel|politics|commerce|economic|religio|census|textbooks|addresses|"
    r"nomenclature|outlines|popular works|miscellanea|research|diseases|control|food|eggs|nests|"
    r"life cycles|migration|abstracts|sources|correspondence|diaries|indexes|exhibitions|museums|"
    r"illustrations|charts|statistics|terminology|psychology|physiology|evolution|variation|"
    r"wintering|outdoor books|homes and haunts|gold discoveries|intellectual life|missions|"
    r"methodology|attributes|roots|seeds|flowers|leaves|fruit|wood|breeding|hunting|trapping|"
    r"conservation|protection|feeding|flight|songs|color|poetry|humor|caricatures|quotations|"
    r"heraldry|genealogy|registers|manners and customs|languages?|races|emigration|"
    r"armed forces|campaigns|colonization|foreign relations|boundaries|geography|climate|"
    r"civilization|ethnic relations|rural conditions|portraits|songs and music)", re.I)
# Single-word LoC subdivisions that Nominatim would otherwise match to some village or peak.
NOISE = {"english", "british", "christianity", "botany", "medicine", "assimilation", "causes",
         "cree", "cretaceous", "zoology", "insects", "birds", "plants", "animals", "science"}
TRAVEL_SUBDIV = re.compile(r"^(description and travel|discovery and exploration)", re.I)
ABBREV = {"Mass.": "Massachusetts", "Me.": "Maine", "N.S.": "Nova Scotia", "S.A.": "South Australia",
          "Colo.": "Colorado", "Calif.": "California", "N.Y.": "New York", "Pa.": "Pennsylvania",
          "Va.": "Virginia", "Fla.": "Florida", "Mich.": "Michigan", "N.H.": "New Hampshire",
          "Vt.": "Vermont", "Conn.": "Connecticut", "Wis.": "Wisconsin", "Minn.": "Minnesota",
          "Ill.": "Illinois", "Ohio": "Ohio", "Ky.": "Kentucky", "Tenn.": "Tennessee",
          "N.C.": "North Carolina", "S.C.": "South Carolina", "Ga.": "Georgia", "La.": "Louisiana",
          "Tex.": "Texas", "Ariz.": "Arizona", "N.M.": "New Mexico", "Utah": "Utah", "Nev.": "Nevada",
          "Or.": "Oregon", "Wash.": "Washington", "Wyo.": "Wyoming", "Mont.": "Montana",
          "B.C.": "British Columbia", "Ont.": "Ontario", "Qué.": "Québec", "N.B.": "New Brunswick",
          "Nfld.": "Newfoundland", "Alta.": "Alberta", "Man.": "Manitoba", "Sask.": "Saskatchewan",
          "N.S.W.": "New South Wales", "Vic.": "Victoria", "Qld.": "Queensland", "Tas.": "Tasmania",
          "W.A.": "Western Australia", "N.Z.": "New Zealand", "Ind.": "Indiana", "Mo.": "Missouri",
          "Md.": "Maryland", "N.J.": "New Jersey", "Del.": "Delaware", "R.I.": "Rhode Island",
          "Ark.": "Arkansas", "Miss.": "Mississippi", "Ala.": "Alabama", "Iowa": "Iowa",
          "Kan.": "Kansas", "Neb.": "Nebraska", "Okla.": "Oklahoma", "Idaho": "Idaho",
          "S.D.": "South Dakota", "N.D.": "North Dakota", "W. Va.": "West Virginia",
          "Hawaii": "Hawaii", "Alaska": "Alaska", "Scot.": "Scotland", "Eng.": "England"}
KEEP_CLASSES = {"boundary", "place", "natural", "waterway", "landuse", "mountain_pass"}
DROP_TYPES = {"house", "building", "hamlet", "isolated_dwelling", "farm", "locality", "neighbourhood",
              "suburb", "quarter", "square", "residential", "cemetery", "grave_yard", "park", "retail"}


def loc_places(subjects):
    """Place headings in a list of LoC subject strings."""
    out = []
    for s in subjects:
        parts = [p.strip() for p in s.split(" -- ")]
        if len(parts) > 1 and TRAVEL_SUBDIV.match(parts[1]):
            out.append(parts[0])
        for p in parts[1:]:
            if not NOT_PLACE.match(p) and not re.search(r"\d", p) and p.lower() not in NOISE:
                out.append(p)
    return out


def loc_query(name):
    """LoC form -> a query Nominatim understands: 'Plymouth (Mass.)' -> 'Plymouth, Massachusetts'."""
    m = re.match(r"^(.*?)\s*\((.*)\)$", name)
    base, qual = (m.group(1), m.group(2)) if m else (name, "")
    m = re.match(r"^(.*), (Mount|Lake|Lakes of|Cape|Gulf of|Bay of|Isle of)$", base)
    if m:
        base = f"{m.group(2)} {m.group(1)}"
    qual = re.sub(r"\b(Extinct city|State|Province|Department|Region|Republic|Ancient)\b", "", qual)
    qual = re.sub(r"\s*:\s*", ", ", qual).strip(" ,:")
    qual = ABBREV.get(qual, qual)
    return f"{base}, {qual}" if qual else base


def nominatim(query, lang):
    """Up to five candidates for a query, or None if Nominatim could not be reached."""
    params = {"q": query, "format": "jsonv2", "limit": 5, "accept-language": lang}
    url = NOMINATIM + "?" + urllib.parse.urlencode(params)
    for attempt in range(3):
        try:
            with urllib.request.urlopen(urllib.request.Request(url, headers=UA), timeout=60) as r:
                hits = json.load(r)
            time.sleep(1.1)  # Nominatim usage policy: at most one request per second
            return [h for h in hits if h.get("category") in KEEP_CLASSES and h.get("type") not in DROP_TYPES]
        except Exception as e:
            print(f"  ! {query}: {e}", file=sys.stderr)
            time.sleep(5 * (attempt + 1))
    return None


def geocode(query, cache, lang):
    """The most important matching place. Taking Nominatim's first hit put "Egypt" in Alabama and
    "Jerusalem, Palestine" on a cemetery there; its importance score ranks the country and the city first."""
    key = f"{lang}|{query}"
    if key in cache:
        return cache[key]
    parts = list(dict.fromkeys(p.strip() for p in query.split(",")))  # "Egypt, Egypt" -> "Egypt"
    hits = nominatim(", ".join(parts), lang)
    if hits is None:
        return None  # not cached, so a later run retries
    best = max(hits, key=lambda h: float(h.get("importance") or 0), default=None)
    if len(parts) > 1 and (best is None or float(best.get("importance") or 0) < 0.35):
        alone = nominatim(parts[0], lang) or []  # "Jerusalem, Palestine" -> "Jerusalem"
        best = max(hits + alone, key=lambda h: float(h.get("importance") or 0), default=None)
    result = None
    if best:
        result = {"lat": round(float(best["lat"]), 3), "lon": round(float(best["lon"]), 3),
                  "name": best.get("name") or query, "type": best.get("addresstype") or best.get("type")}
    cache[key] = result
    return result


def llm_places(texts, cache_path, ask=True, workers=4):
    """Map each title to a list of English place names, 40 titles per request."""
    cache = json.loads(Path(cache_path).read_text(encoding="utf-8")) if Path(cache_path).exists() else {}
    key = lambda t: hashlib.sha1(t.encode()).hexdigest()[:16]
    todo = [t for t in dict.fromkeys(texts) if key(t) not in cache] if ask else []
    batches = [todo[i:i + 40] for i in range(0, len(todo), 40)]

    def run(batch):
        prompt = LLM_PROMPT + "\n".join(f"{n}. {t}" for n, t in enumerate(batch, 1))
        body = {"model": LLM_MODEL, "messages": [{"role": "user", "content": prompt}],
                "temperature": 0, "max_tokens": 4000,
                "chat_template_kwargs": {"enable_thinking": False}}
        req = urllib.request.Request(INNKUBE, data=json.dumps(body).encode(), headers={
            "Authorization": "Bearer " + os.environ["INNKUBE_API_KEY"], "Content-Type": "application/json"})
        try:
            with urllib.request.urlopen(req, timeout=300) as r:
                content = json.load(r)["choices"][0]["message"]["content"]
            content = re.sub(r"<think>.*?</think>|```(json)?", "", content, flags=re.S)
            return batch, json.loads(content[content.index("{"):content.rindex("}") + 1])
        except Exception as e:  # leave the batch uncached; a rerun retries it
            print(f"  ! LLM batch: {e}", file=sys.stderr)
            return batch, None

    with ThreadPoolExecutor(workers) as pool:
        for i, (batch, answer) in enumerate(pool.map(run, batches), 1):
            if answer is None:
                continue
            for n, t in enumerate(batch, 1):
                got = answer.get(str(n), [])
                cache[key(t)] = [p for p in got if isinstance(p, str) and p.strip()]
            Path(cache_path).write_text(json.dumps(cache, ensure_ascii=False), encoding="utf-8")
            print(f"  LLM {i}/{len(batches)}", file=sys.stderr)
    return {t: cache.get(key(t), []) for t in texts}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", default="overview/data.json")
    ap.add_argument("--places", default="overview/places.json")
    ap.add_argument("--cache", default="geocache.json")
    ap.add_argument("--llm-cache", default="llm_places.json")
    ap.add_argument("--cached-only", action="store_true",
                    help="don't ask Nominatim or the LLM; use only answers already cached (for a quick partial build)")
    ap.add_argument("--llm-only", action="store_true", help="only fill the LLM cache; don't geocode or write data")
    ap.add_argument("--recheck-americas", action="store_true",
                    help="geocode again every cached name that resolved to the Americas (after the importance fix)")
    args = ap.parse_args()

    rows = json.load(open(args.data, encoding="utf-8"))
    cache = json.loads(Path(args.cache).read_text(encoding="utf-8")) if Path(args.cache).exists() else {}
    if args.recheck_americas:
        stale = [k for k, v in cache.items() if v and -170 < v["lon"] < -30 and -56 < v["lat"] < 72]
        for k in stale:
            del cache[k]
        print(f"{len(stale)} names placed in the Americas will be geocoded again", file=sys.stderr)

    # 1. Names per work
    names = {}
    for r in rows:
        if r["src"] == "PG":
            names[id(r)] = [("en", p) for p in loc_places((r.get("sub") or "").split(" | ") if r.get("sub") else [])]
    by_llm = [r for r in rows if r["src"] != "PG" or not names[id(r)]]
    text = lambda r: " — ".join(filter(None, [r["t"], r.get("sum", "") or (r.get("sub", "") if r["src"] in ("BHL", "IA", "GA", "BNE", "UCM", "BNP") else "")]))
    found = llm_places([text(r) for r in by_llm], args.llm_cache, ask=not args.cached_only)
    if args.llm_only:
        return
    for r in by_llm:
        names[id(r)] = [("en", p.strip()) for p in found[text(r)]]

    # 2. Geocode each distinct name once
    distinct = Counter(n for ns in names.values() for n in set(ns))
    print(f"{len(distinct)} distinct place names", file=sys.stderr)
    places, resolved = {}, {}
    todo = [n for n in distinct if n[1].lower() not in REGIONS
            and f"{n[0]}|{loc_query(n[1])}" not in cache]
    print(f"{len(todo)} need Nominatim (~{len(todo) * 1.2 / 60:.0f} min)", file=sys.stderr)
    for i, (lang, name) in enumerate(distinct):
        if name.lower() in REGIONS:
            lat, lon = REGIONS[name.lower()]
            g = {"lat": lat, "lon": lon, "name": name, "type": "region"}
        elif args.cached_only:
            g = cache.get(f"{lang}|{loc_query(name)}")
        else:
            g = geocode(loc_query(name), cache, lang)
        resolved[(lang, name)] = g
        if i % 50 == 0 and not args.cached_only:
            Path(args.cache).write_text(json.dumps(cache, ensure_ascii=False), encoding="utf-8")
            print(f"  {i}/{len(distinct)}", file=sys.stderr)
    if not args.cached_only:
        Path(args.cache).write_text(json.dumps(cache, ensure_ascii=False), encoding="utf-8")

    # 3. Merge names that resolve to the same point; attach place keys to works
    for r in rows:
        keys = []
        for n in names.get(id(r), []):
            g = resolved.get(n)
            if not g:
                continue
            key = f"{g['lat']:.2f},{g['lon']:.2f}"
            p = places.setdefault(key, {"lat": g["lat"], "lon": g["lon"], "labels": Counter(), "type": g["type"]})
            p["labels"][n[1]] += 1
            if key not in keys:
                keys.append(key)
        if keys:
            r["p"] = keys
        else:
            r.pop("p", None)
    for p in places.values():
        p["label"] = p.pop("labels").most_common(1)[0][0]

    json.dump(rows, open(args.data, "w", encoding="utf-8"), ensure_ascii=False, separators=(",", ":"))
    json.dump(places, open(args.places, "w", encoding="utf-8"), ensure_ascii=False, separators=(",", ":"))
    by_src = Counter(r["src"] for r in rows if r.get("p"))
    tot = Counter(r["src"] for r in rows)
    print(f"Wrote {len(places)} places to {args.places}", file=sys.stderr)
    for s in tot:
        print(f"  {s}: {by_src[s]}/{tot[s]} works located", file=sys.stderr)


if __name__ == "__main__":
    main()