#!/usr/bin/env python3
"""
Give every work in overview/data.json a genre, and drop works that are not natural history or
travel writing.

Usage:
    python genres.py               # after build_overview.py, before years.py

The harvesters' first pass asked the LLM only "natural history, travel, both or neither", and
defined travel loosely enough ("a description of a country or region as seen by a traveller")
that the Internet Archive and Gallica, whose broadest subject heading is "Description and
travel", brought in guidebooks, county topographies, local histories, emigration tracts,
devotional books and novels. This pass asks the LLM on innkube (qwen38-27b, Uni Passau) for one
genre per work from the list below. A kept genre adds its category ("c") to the ones the work
already had (a voyage that was natural history stays natural history too). Works in the other
genres are dropped, with two exceptions: the hand-curated sources (DTA's fixed genres, Humboldt's
journals and letters, Darwin and Wallace, EcoCor) keep their works, and natural-history works,
which came in through precise subject headings, are dropped only as fiction, devotional books,
promotion or guidebooks; as technical science, topography, history or periodicals they stay
natural history and show that genre (periodicals are flagged and hidden by default on the page).
The genre goes into "gn". Answers are cached in genre_labels.json. Needs INNKUBE_API_KEY.
"""
import argparse
import json
import os
import re
import sys
import urllib.request
from collections import Counter
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

INNKUBE = "https://llms.innkube.fim.uni-passau.de/v1/chat/completions"
LLM_MODEL = "qwen38-27b"
BATCH = 40
# (key, label on the page, description for the LLM, categories, kept)
GENRES = [
    ("travel_account", "Travel account", "a first-hand account of a journey, tour, pilgrimage or stay abroad, "
     "including travel letters and diaries, and memoirs of years spent abroad or on a frontier", ["travel"], True),
    ("voyage", "Voyage or expedition", "a sea voyage, circumnavigation, or an exploring, surveying or scientific "
     "expedition, told by a participant, and collections of such voyages", ["travel"], True),
    ("sketches", "Sketches of a region", "impressions, rambles or sketches of a country, region or people by a "
     "visitor or observant resident (not a guidebook or a topography)", ["travel"], True),
    ("sport", "Hunting, sport and mountaineering", "accounts of big-game hunting, fishing, shooting, climbing and "
     "mountaineering", ["travel"], True),
    ("naturalist_travels", "Naturalist's travels", "a naturalist's journey, voyage or stay with substantial "
     "observation of animals, plants or landscape", ["natural_history", "travel"], True),
    ("natural_history", "Natural history", "descriptive natural history of animals, plants, minerals or a region, "
     "systematic or popular (Buffon, Brehm, regional natural histories), and works on evolution, distribution and "
     "the order of nature for learned general readers (Darwin's Origin, Humboldt's Kosmos)", ["natural_history"], True),
    ("nature_essays", "Nature essays and observation", "essays, observations and rambles about nature, popular "
     "nature study, natural history in a literary vein (Gilbert White, Thoreau)", ["natural_history"], True),
    ("letters_journals", "Letters and journals", "correspondence, diaries and working journals of naturalists or "
     "travellers that are not one journey's account", None, True),
    ("guidebook", "Guidebook", "guides, itineraries, handbooks and route books for travellers or tourists", [], False),
    ("topography", "Topography or local history", "gazetteers, statistical or topographical descriptions, county "
     "and city descriptions, local and regional histories, picture books and collections of views", [], False),
    ("promotion", "Emigration or promotion", "emigrants' handbooks, settlement, railway, resort or commercial "
     "promotion", [], False),
    ("devotional", "Devotional", "devotional books, sermons, Bible geography and religious instruction", [], False),
    ("science", "Technical science", "technical or specialist scientific works: taxonomic monographs, species "
     "lists, reports, proceedings, textbooks, manuals", [], False),
    ("history", "History or biography", "histories (including later histories of exploration), biographies, "
     "war diaries, ethnographic or political treatises", [], False),
    ("fiction", "Fiction or poetry", "novels, stories, poems, plays and juvenile fiction", [], False),
    ("periodical", "Periodical", "journals, magazines and serial publications as a whole", [], False),
    ("other", "Other", "anything else", [], False),
]
BY_KEY = {g[0]: g for g in GENRES}
CURATED = {"DTA", "EHD", "DWO", "EC"}
# Natural history came in through precise subject headings (natural history, zoology, botany ...), unlike
# travel's broad "Description and travel", so a natural-history work is only dropped for these genres.
NH_KEEPS = {"science", "topography", "history", "periodical", "other"}
SRC_NAME = {"PG": "Project Gutenberg", "BHL": "Biodiversity Heritage Library", "IA": "Internet Archive",
            "GA": "Gallica", "DTA": "Deutsches Textarchiv", "WS": "Wikisource", "DWO": "Darwin Online",
            "EHD": "edition humboldt digital", "EC": "EcoCor", "BNE": "Biblioteca Nacional de España",
            "UCM": "Universidad Complutense de Madrid", "BNP": "Biblioteca Nacional de Portugal",
}
PROMPT = ("Each numbered line is a work in a catalogue of historical natural histories and travel writing: "
          "title | author | year | library | subject headings, genre or description. Give each work exactly one "
          "genre:\n" + "\n".join(f"- {k}: {d}" for k, _, d, _, _ in GENRES)
          + "\nJudge by what the work is, not by what it is about: a guidebook to the Alps is a guidebook, a "
            "climber's account of the Alps is sport, and a traveller's tour that also surveys a country's resources is "
            "still a travel account. Answer only with JSON mapping each number to one genre key.\n\n")


def describe(r):
    ctx = r.get("sub") or r.get("sum") or r.get("g") or ""
    if r["src"] == "PG":
        ctx = r.get("sub") or ""
    return (f"{r['t'][:200]} | {(r.get('a') or 'anonymous')[:60]} | {r.get('y') or ''} | {SRC_NAME[r['src']]} | "
            f"{ctx[:220]}")


def classify(rows, cache_path, workers):
    cache = json.loads(Path(cache_path).read_text(encoding="utf-8")) if Path(cache_path).exists() else {}
    key = lambda r: f"{r['src']}:{r['id']}"
    todo = [r for r in rows if key(r) not in cache]
    batches = [todo[i:i + BATCH] for i in range(0, len(todo), BATCH)]
    print(f"{len(todo)} works to sort into genres ({len(batches)} requests)", file=sys.stderr)

    def ask(batch):
        prompt = PROMPT + "\n".join(f"{n}. {describe(r)}" for n, r in enumerate(batch, 1))
        body = {"model": LLM_MODEL, "messages": [{"role": "user", "content": prompt}], "temperature": 0,
                "max_tokens": 3000, "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 resp:
                content = json.load(resp)["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 uncached; a rerun retries
            print(f"  ! batch failed: {e}", file=sys.stderr)
            return batch, None

    with ThreadPoolExecutor(workers) as pool:
        for n, (batch, answer) in enumerate(pool.map(ask, batches), 1):
            if answer:
                for i, r in enumerate(batch, 1):
                    if answer.get(str(i)) in BY_KEY:
                        cache[key(r)] = answer[str(i)]
            if n % 10 == 0 or n == len(batches):
                Path(cache_path).write_text(json.dumps(cache, ensure_ascii=False), encoding="utf-8")
                print(f"  {n}/{len(batches)}", file=sys.stderr)
    return {key(r): cache.get(key(r)) for r in rows}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", default="overview/data.json")
    ap.add_argument("--cache", default="genre_labels.json")
    ap.add_argument("--workers", type=int, default=6)
    args = ap.parse_args()
    rows = json.load(open(args.data, encoding="utf-8"))
    genres = classify(rows, args.cache, args.workers)

    out, dropped, unlabelled = [], Counter(), 0
    for r in rows:
        g = genres[f"{r['src']}:{r['id']}"]
        if g is None:  # the LLM gave no answer: keep the work as it was; a rerun asks again
            unlabelled += 1
            out.append(r)
            continue
        _, _, _, cats, kept = BY_KEY[g]
        was_nh = "natural_history" in r["c"]
        if not kept and r["src"] not in CURATED and not (was_nh and g in NH_KEEPS):
            dropped[(r["src"], g)] += 1
            continue
        r["gn"] = g
        if kept:  # a genre adds its categories but never takes one away
            r["c"] = [c for c in ("natural_history", "travel") if c in r["c"] or c in (cats or [])]
        elif was_nh:  # natural history kept despite its genre, e.g. a regional description of a coast's flora
            r["c"] = ["natural_history"]
        if g == "periodical" and "periodical" not in r["f"]:
            r["f"] = r["f"] + ["periodical"]  # hidden by default on the page, like Gutenberg's periodicals
        out.append(r)
    json.dump(out, open(args.data, "w", encoding="utf-8"), ensure_ascii=False, separators=(",", ":"))
    print(f"{len(rows)} works -> {len(out)} kept; {sum(dropped.values())} dropped"
          + (f"; {unlabelled} without an answer (kept; rerun to retry)" if unlabelled else ""), file=sys.stderr)
    by_src = Counter()
    for (src, g), n in dropped.items():
        by_src[src] += n
    for src, n in by_src.most_common():
        detail = ", ".join(f"{g} {m}" for (s, g), m in sorted(dropped.items(), key=lambda x: -x[1]) if s == src)
        print(f"  {src}: {n} dropped ({detail})", file=sys.stderr)
    print("  kept by genre:", dict(Counter(r.get("gn") for r in out).most_common()), file=sys.stderr)


if __name__ == "__main__":
    main()