#!/usr/bin/env python3
"""WP44 — select the synthesis populations (protocol P2/S5):
  high-priority set (target >= 10,000): per (primary category × publication year) the top works by NCS36 (field-normalised
  36-month citations; for 2025–2026 by c12 then cited_by_count), proportional allocation with a floor of 25 per cell, eligible
  works only, retracted excluded;
  landmark set (target >= 1,000): top by NCS36 within category (floor 15 per category) plus the top 200 overall by c36,
  deduplicated; flagged for PI review (H-3).
Outputs: literature_synthesis/landmarks/{high_priority.parquet,landmarks.csv,<area>_landmarks.csv,<area>_abstracts.jsonl,selection_report.json}."""
import argparse, json, os, duckdb, yaml, pandas as pd


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--root", required=True); ap.add_argument("--ontology", required=True)
    ap.add_argument("--n-high", type=int, default=10000); ap.add_argument("--n-land", type=int, default=1000); a = ap.parse_args(); R = a.root
    out = f"{R}/literature_synthesis/landmarks"; os.makedirs(out, exist_ok=True)
    ont = yaml.safe_load(open(a.ontology)); area = {c["id"]: c["area"] for c in ont["categories"]}; cats = list(area)
    con = duckdb.connect(f"{R}/data/processed/wp44.duckdb", read_only=True); con.execute("SET threads TO 16;")
    WP = f"read_parquet('{R}/data/processed/classification/primary/*.parquet')"; WCW = f"read_parquet('{R}/data/processed/corpus/work_citation_windows.parquet')"
    df = con.execute(f"""SELECT c.id AS work_id, c.doi, c.arxiv_id_linked AS arxiv_id, c.title, c.source_name, c.publication_year AS year, c.type, p.primary_category, p.n_categories,
                               w.c12, w.c36, w.ncs36, c.cited_by_count, c.fwci, c.has_abstract, c.is_oa, c.period
                        FROM corpus_works c JOIN {WP} p ON p.work_id = c.id JOIN {WCW} w ON w.work_id = c.id WHERE c.eligible AND NOT c.e6_retracted""").df()
    # cleaning (2026-11-03, after Agent H's audit of the first list): drop venue-like records (title equal to the source name
    # or shorter than four words), records with neither abstract nor reference list, and duplicate titles (keep the earliest
    # year, then the most cited) — OpenAlex carries SSRN/reprint duplicates of landmark papers
    import re as _re
    df["ntitle"] = df.title.fillna("").str.lower().str.replace(r"[^a-z0-9]+", "", regex=True)
    df = df[(df.title.fillna("").str.split().str.len() >= 4) & (df.title.fillna("").str.lower() != df.source_name.fillna("").str.lower())]
    df = df[df.has_abstract | (df.cited_by_count > 0)]
    df = df.sort_values(["ntitle", "year", "c36"], ascending=[True, True, False]).drop_duplicates("ntitle", keep="first")
    df["score"] = df.ncs36.fillna(0); recent = df.year >= 2025; df.loc[recent, "score"] = df.loc[recent, "c12"].fillna(0) / (df.loc[recent].groupby("year").c12.transform("mean") + 1e-9)
    cell = df.groupby(["primary_category", "year"]).size(); alloc = (cell / cell.sum() * a.n_high).clip(lower=25).round().astype(int)
    parts = []
    for (c, y), k in alloc.items():
        g = df[(df.primary_category == c) & (df.year == y)].nlargest(int(k), ["score", "cited_by_count"]); parts.append(g)
    high = pd.concat(parts).drop_duplicates("work_id"); high["set"] = "high_priority"
    land_parts = [df[df.primary_category == c].nlargest(max(15, a.n_land // len(cats)), ["score", "cited_by_count"]) for c in cats]
    land = pd.concat(land_parts + [df.nlargest(200, ["c36", "cited_by_count"])]).drop_duplicates("work_id"); land["set"] = "landmark"; land["pi_review"] = True
    high.to_parquet(f"{out}/high_priority.parquet"); land.to_csv(f"{out}/landmarks.csv", index=False)
    L = f"read_csv_auto('{out}/landmarks.csv')"
    for ar in sorted(set(area.values())):
        ids = [c for c in cats if area[c] == ar]; sub = land[land.primary_category.isin(ids)]
        sub.to_csv(f"{out}/{ar}_landmarks.csv", index=False)
        ab = con.execute(f"SELECT l.work_id, l.title, a.abstract FROM {L} l LEFT JOIN oa_abstracts a ON a.work_id = l.work_id WHERE l.primary_category IN ({','.join(repr(i) for i in ids)})").fetchall()
        with open(f"{out}/{ar}_abstracts.jsonl", "w") as fh:
            for w, t, abst in ab:
                fh.write(json.dumps({"work_id": int(w), "title": t, "abstract": abst}) + "\n")
    rep = {"high_priority": int(len(high)), "landmarks": int(len(land)), "by_area": land.primary_category.map(area).value_counts().to_dict(), "by_year_landmarks": land.year.value_counts().sort_index().to_dict(),
           "share_with_abstract_landmarks": float(land.has_abstract.mean()), "share_oa_landmarks": float(land.is_oa.mean())}
    json.dump(rep, open(f"{out}/selection_report.json", "w"), indent=1); print(rep)


if __name__ == "__main__":
    main()