#!/usr/bin/env python3
"""WP44 — Phase 2: assign the 40 ontology categories (multi-label) to every work of the corpus.

Signals per work: OpenAlex topics (ids + subfield ids), arXiv categories (through link_arxiv_oa), DBLP venue stream
(through link_dblp_oa), keyword regexes on title and abstract. Scoring follows ontology.yaml `assignment_rule`:
per signal type the matched weights are summed and capped at 1.0, then summed over signal types; a category is assigned
if score >= threshold; primary = argmax (ties: child before parent, then lower category number).

Step 1 (`--prepare`): DuckDB builds data/processed/classification/input/chunk_*.parquet (id, year, title, abstract,
topic_ids, subfield_ids, arxiv_cats, dblp_streams). Step 2 (`--run`): multiprocessing over chunks →
data/processed/classification/work_categories.parquet (work_id, category_id, score, s_topic, s_subfield, s_arxiv,
s_dblp, s_kw) and work_primary.parquet (work_id, primary_category, n_categories). Deterministic; version-stamped.
"""
import argparse, glob, json, multiprocessing as mp, os, re, time, datetime as dt
import yaml, duckdb
import pyarrow as pa, pyarrow.parquet as pq

ONT = None


def load_ontology(path):
    o = yaml.safe_load(open(path))
    cats = o["categories"]
    comp = []
    for c in cats:
        kws = []
        for k in (c.get("keywords") or []):
            kws.append((re.compile(k["pattern"]), float(k["weight"]), k.get("field", "both")))
        exc = [re.compile(p if isinstance(p, str) else p["pattern"]) for p in (c.get("exclusions") or [])]
        comp.append(dict(
            id=c["id"], number=int(c["number"]), parent=tuple(c["child_of"]) if isinstance(c.get("child_of"), list) else ((c["child_of"],) if c.get("child_of") else ()),
            topics={str(t["id"]).lstrip("T"): float(t["weight"]) for t in (c.get("openalex_topics") or [])},
            subfields={str(t["id"]): float(t["weight"]) for t in (c.get("openalex_subfields") or [])},
            arxiv={t["id"]: float(t["weight"]) for t in (c.get("arxiv_categories") or [])},
            dblp={t["key"]: float(t["weight"]) for t in (c.get("dblp_streams") or [])},
            kws=kws, exc=exc))
    rule = o.get("assignment_rule", {})
    return dict(cats=comp, threshold=float(rule.get("threshold", 1.0)), version=o.get("version"))


def _init(path):
    global ONT
    ONT = load_ontology(path)


def score_row(title, abstract, topic_ids, subfield_ids, arxiv_cats, dblp_streams):
    title = title or ""; abstract = abstract or ""; both = title + "\n" + abstract
    out = []
    for c in ONT["cats"]:
        st = min(1.0, sum(c["topics"].get(str(t), 0.0) for t in set(topic_ids or [])))
        ss = min(1.0, sum(c["subfields"].get(str(s), 0.0) for s in set(subfield_ids or [])))
        sa = min(1.0, sum(c["arxiv"].get(a, 0.0) for a in set(arxiv_cats or [])))
        sd = min(1.0, sum(c["dblp"].get(d, 0.0) for d in set(dblp_streams or [])))
        sk = 0.0
        for rx, w, field in c["kws"]:
            txt = title if field == "title" else abstract if field == "abstract" else both
            if txt and rx.search(txt):
                sk += w
                if sk >= 1.0:
                    sk = 1.0; break
        s = st + ss + sa + sd + sk
        if s >= ONT["threshold"] and c["exc"] and any(rx.search(both) for rx in c["exc"]):
            s = 0.0
        if s >= ONT["threshold"]:
            out.append((c["id"], c["number"], c["parent"], round(s, 3), st, ss, sa, sd, sk))
    return out


def process_chunk(args):
    path, out_dir = args
    name = os.path.basename(path).replace(".parquet", "")
    o1 = os.path.join(out_dir, "cats", name + ".parquet"); o2 = os.path.join(out_dir, "primary", name + ".parquet")
    if os.path.exists(o1) and os.path.exists(o2):
        return name, 0, 0
    t = pq.read_table(path).to_pydict()
    cw, cc, cs, c1, c2, c3, c4, c5 = [], [], [], [], [], [], [], []
    pw, pc, pn, ps = [], [], [], []
    n = len(t["id"])
    for i in range(n):
        res = score_row(t["title"][i], t["abstract"][i], t["topic_ids"][i], t["subfield_ids"][i], t["arxiv_cats"][i], t["dblp_streams"][i])
        for (cid, num, parent, s, st, ss, sa, sd, sk) in res:
            cw.append(t["id"][i]); cc.append(cid); cs.append(s); c1.append(st); c2.append(ss); c3.append(sa); c4.append(sd); c5.append(sk)
        if res:
            parents = set().union(*[set(r[2]) for r in res])
            # primary: highest score; tie → child before parent (a category that is someone's parent loses), then lower number
            best = sorted(res, key=lambda r: (-r[3], 1 if r[0] in parents else 0, r[1]))[0]
            pw.append(t["id"][i]); pc.append(best[0]); pn.append(len(res)); ps.append(best[3])
    os.makedirs(os.path.dirname(o1), exist_ok=True); os.makedirs(os.path.dirname(o2), exist_ok=True)
    pq.write_table(pa.table({"work_id": pa.array(cw, pa.int64()), "category_id": cc, "score": pa.array(cs, pa.float32()), "s_topic": pa.array(c1, pa.float32()),
                             "s_subfield": pa.array(c2, pa.float32()), "s_arxiv": pa.array(c3, pa.float32()), "s_dblp": pa.array(c4, pa.float32()), "s_kw": pa.array(c5, pa.float32())}), o1 + ".part", compression="zstd")
    os.replace(o1 + ".part", o1)
    pq.write_table(pa.table({"work_id": pa.array(pw, pa.int64()), "primary_category": pc, "n_categories": pa.array(pn, pa.int16()), "primary_score": pa.array(ps, pa.float32())}), o2 + ".part", compression="zstd")
    os.replace(o2 + ".part", o2)
    return name, n, len(pw)


def prepare(root, db, out_dir, chunk_rows):
    con = duckdb.connect(db); con.execute("SET threads TO 24; SET memory_limit='" + os.environ.get("WP44_DUCKDB_MEM", "60GB") + "';")
    os.makedirs(os.path.join(out_dir, "input"), exist_ok=True)
    con.execute("""CREATE OR REPLACE TABLE cls_input AS
        SELECT w.id, w.publication_year AS year, w.title, a.abstract,
               (SELECT list(topic_id ORDER BY rank) FROM oa_work_topics t WHERE t.work_id = w.id) AS topic_ids,
               (SELECT list(subfield_id ORDER BY rank) FROM oa_work_topics t WHERE t.work_id = w.id) AS subfield_ids,
               (SELECT flatten(list(x.categories)) FROM link_arxiv_oa l JOIN arxiv_keys x USING(arxiv_id) WHERE l.oa_id = w.id) AS arxiv_cats,
               (SELECT list(DISTINCT d.stream) FROM link_dblp_oa l JOIN dblp_keys d USING(dblp_key) WHERE l.oa_id = w.id) AS dblp_streams
        FROM corpus_works w LEFT JOIN oa_abstracts a ON a.work_id = w.id WHERE w.eligible OR w.cs_penumbra""")
    n = con.execute("SELECT count(*) FROM cls_input").fetchone()[0]
    con.execute(f"""COPY (SELECT *, (row_number() OVER ()) // {chunk_rows} AS chunk FROM cls_input) TO '{out_dir}/input' (FORMAT PARQUET, PARTITION_BY (chunk), OVERWRITE_OR_IGNORE, COMPRESSION ZSTD)""")
    print("prepared", n, "works")


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--root", required=True); ap.add_argument("--ontology", required=True)
    ap.add_argument("--prepare", action="store_true"); ap.add_argument("--run", action="store_true"); ap.add_argument("--workers", type=int, default=20)
    ap.add_argument("--chunk-rows", type=int, default=100000); ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--shard", default=None, help="i/N: process only chunks with index %% N == i in this process (no pool; used by run_classify_shards.sh)")
    a = ap.parse_args(); out_dir = os.path.join(a.root, "data/processed/classification"); db = os.path.join(a.root, "data/processed/wp44.duckdb")
    if a.prepare:
        prepare(a.root, db, out_dir, a.chunk_rows)
    if a.run and a.shard:
        i, N = map(int, a.shard.split("/")); _init(a.ontology)
        files = sorted(glob.glob(os.path.join(out_dir, "input", "chunk=*", "*.parquet")), key=lambda f: int(os.path.basename(os.path.dirname(f)).split("=")[1]))
        t0 = time.time(); tot = cl = 0
        for k, f in enumerate(files):
            if k % N != i:
                continue
            name, n, c = process_chunk_named((f, out_dir, os.path.basename(os.path.dirname(f)))); tot += n; cl += c
            print(f"[{time.time()-t0:6.0f}s] shard {i}/{N} {name}: {n} works, {c} classified | {tot} / {cl}", flush=True)
        return
    if a.run:
        files = sorted(glob.glob(os.path.join(out_dir, "input", "chunk=*", "*.parquet")))
        files = [(f, out_dir) for f in files]
        if a.limit:
            files = files[: a.limit]
        # chunk files share the name data_0.parquet inside chunk=N dirs → rename key
        files = [(f, out_dir) for f, _ in files]
        t0 = time.time(); tot = 0; cl = 0
        with mp.Pool(a.workers, initializer=_init, initargs=(a.ontology,)) as pool:
            for name, n, k in pool.imap_unordered(process_chunk_named, [(f, out_dir, os.path.basename(os.path.dirname(f))) for f, _ in files]):
                tot += n; cl += k
                print(f"[{time.time()-t0:6.0f}s] {name}: {n} works, {k} classified | total {tot} / classified {cl}", flush=True)
        json.dump({"ontology": a.ontology, "version": load_ontology(a.ontology)["version"], "works": tot, "classified": cl, "when": dt.datetime.now(dt.UTC).isoformat(),
                   "seconds": round(time.time() - t0)}, open(os.path.join(out_dir, "classification_report.json"), "w"), indent=1)


def process_chunk_named(args):
    path, out_dir, chunk_name = args
    name = chunk_name.replace("=", "_")
    o1 = os.path.join(out_dir, "cats", name + ".parquet"); o2 = os.path.join(out_dir, "primary", name + ".parquet")
    if os.path.exists(o1) and os.path.exists(o2):
        return name, 0, 0
    # reuse process_chunk with a symlinked name
    tmp = os.path.join(out_dir, "input_named"); os.makedirs(tmp, exist_ok=True)
    link = os.path.join(tmp, name + ".parquet")
    if not os.path.exists(link):
        os.symlink(path, link)
    return process_chunk((link, out_dir))


if __name__ == "__main__":
    main()