#!/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.
- DTA, Wikisource and EcoCor: an LLM on innkube (qwen38-27b, Uni Passau) reads the title and,
  for Wikisource, the short description, 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 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 (mostly German) 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 = {
    "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),
}
# 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"}


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 geocode(query, cache, lang):
    key = f"{lang}|{query}"
    if key in cache:
        return cache[key]
    params = {"q": query, "format": "jsonv2", "limit": 1, "accept-language": lang}
    url = NOMINATIM + "?" + urllib.parse.urlencode(params)
    result = None
    for attempt in range(3):
        try:
            with urllib.request.urlopen(urllib.request.Request(url, headers=UA), timeout=60) as r:
                hits = json.load(r)
            break
        except Exception as e:
            print(f"  ! {query}: {e}", file=sys.stderr)
            time.sleep(5 * (attempt + 1))
    else:
        return None  # not cached, so a later run retries
    time.sleep(1.1)  # Nominatim usage policy: at most one request per second
    if hits:
        h = hits[0]
        if h.get("category") in KEEP_CLASSES and h.get("type") not in DROP_TYPES:
            result = {"lat": round(float(h["lat"]), 3), "lon": round(float(h["lon"]), 3),
                      "name": h.get("name") or query, "type": h.get("addresstype") or h.get("type")}
    cache[key] = result
    return result


def llm_places(texts, cache_path):
    """Map each German text to a list of English place names, 40 texts 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]
    for i in range(0, len(todo), 40):
        batch = todo[i:i + 40]
        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)
            answer = json.loads(content[content.index("{"):content.rindex("}") + 1])
        except Exception as e:  # leave the batch uncached; a rerun retries it
            print(f"  ! LLM batch {i // 40}: {e}", file=sys.stderr)
            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 {min(i + 40, len(todo))}/{len(todo)}", 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; use only names already in the cache (for a quick partial build)")
    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 {}

    # 1. Names per work
    names = {}
    german = [r for r in rows if r["src"] in ("DTA", "WS", "EC")]
    text = lambda r: " — ".join(filter(None, [r["t"], r.get("sum", "")]))
    found = llm_places([text(r) for r in german], args.llm_cache)
    for r in german:
        names[id(r)] = [("en", p.strip()) for p in found[text(r)]]
    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 [])]

    # 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()