"""Fetch the citation-year histograms the sealed rule asks for.

Writes three files into ~/data/retraction-citations/:
  sample.csv    the 400 retracted papers, as drawn
  cites.csv     one row per (work, citation year, count), both arms
  controls.csv  the control works and which retracted paper they were drawn for

Nothing here computes a ratio. Analysis lives in analyse.py so that a rerun of
the fetch cannot quietly change a number in the entry.
"""

import csv, json, os, random, sys, time
import urllib.parse, urllib.request
from concurrent.futures import ThreadPoolExecutor

MAIL = "vesper@untilnextsession.com"
API = "https://api.openalex.org/works"
OUT = os.path.expanduser("~/data/retraction-citations")
RW = os.path.expanduser("~/data/retractionwatch-2026-09-20.csv")
SEED = 20260920
TARGET = 400

# OpenAlex 429'd the first run of this script: eight threads with no pause
# between requests is not the polite pool, whatever the mailto says. One
# request every 0.3 s, three at a time, and a long sit-down on a 429.
_threading = __import__("threading")
_throttle = _threading.Semaphore(3)
_pace = _threading.Lock()
_last = [0.0]


def _wait_turn():
    with _pace:
        gap = time.time() - _last[0]
        if gap < 0.30:
            time.sleep(0.30 - gap)
        _last[0] = time.time()


def get(url, tries=5):
    for attempt in range(tries):
        with _throttle:
            _wait_turn()
            try:
                req = urllib.request.Request(
                    url, headers={"User-Agent": "vesper/1.0 (%s)" % MAIL})
                with urllib.request.urlopen(req, timeout=40) as r:
                    return json.loads(r.read().decode())
            except urllib.error.HTTPError as exc:
                if exc.code == 429:
                    time.sleep(20)
                    continue
                if attempt == tries - 1:
                    print("  give up:", url[:110], exc, file=sys.stderr)
                    return None
                time.sleep(1.5 * (attempt + 1))
            except Exception as exc:
                if attempt == tries - 1:
                    print("  give up:", url[:110], exc, file=sys.stderr)
                    return None
                time.sleep(1.5 * (attempt + 1))
    return None


def year_of(value):
    """Retraction Watch dates are `M/D/YYYY H:MM`. Year only is all we use."""
    try:
        return int(str(value).split("/")[2].split(" ")[0])
    except Exception:
        return None


def population():
    """Every row the sealed rule admits, before the random draw."""
    rows, reinstated = [], set()
    with open(RW, newline="", encoding="utf-8", errors="replace") as fh:
        for row in csv.DictReader(fh):
            if row["RetractionNature"] == "Reinstatement":
                reinstated.add(row["Record ID"])
            rows.append(row)
    kept, seen = [], set()
    for row in rows:
        if row["RetractionNature"] != "Retraction":
            continue
        if row["Record ID"] in reinstated:
            continue
        r, p = year_of(row["RetractionDate"]), year_of(row["OriginalPaperDate"])
        if r is None or p is None or not (2018 <= r <= 2021) or p > r - 4:
            continue
        doi = (row["OriginalPaperDOI"] or "").strip().lower()
        # A comma or a pipe in a DOI would be read as filter syntax by OpenAlex.
        if not doi.startswith("10.") or "," in doi or "|" in doi:
            continue
        if doi in seen:
            continue
        seen.add(doi)
        kept.append({"record_id": row["Record ID"], "doi": doi,
                     "retraction_year": r, "paper_year": p,
                     "subject": row["Subject"], "journal": row["Journal"]})
    return kept


def resolve(batch):
    """DOI -> OpenAlex id, publication year, field. 50 at a time."""
    q = "|".join(b["doi"] for b in batch)
    url = "%s?filter=doi:%s&per-page=50&mailto=%s&select=id,doi,publication_year,primary_topic,cited_by_count,is_retracted" % (
        API, urllib.parse.quote(q, safe="|/."), MAIL)
    data = get(url)
    found = {}
    for w in (data or {}).get("results", []):
        doi = (w.get("doi") or "").lower().replace("https://doi.org/", "")
        topic = w.get("primary_topic") or {}
        field = ((topic.get("field") or {}).get("id") or "").split("/")[-1]
        found[doi] = {"openalex": w["id"].split("/")[-1],
                      "oa_year": w.get("publication_year"),
                      "field": field,
                      "lifetime_cites": w.get("cited_by_count"),
                      "oa_retracted": bool(w.get("is_retracted"))}
    return found


def histogram(work_id):
    url = "%s?filter=cites:%s&group_by=publication_year&mailto=%s" % (
        API, work_id, MAIL)
    data = get(url)
    if data is None:
        return None
    return {int(g["key"]): g["count"] for g in data.get("group_by", [])
            if str(g["key"]).isdigit()}


def controls_for(paper):
    """Three random same-year same-field works, never retracted.

    `cited_by_count:>4` is the enrichment the rule declares: a uniform sample
    of a publication year is almost entirely works with no citations, which
    would populate none of the strata the retracted papers occupy.
    """
    if not paper.get("field") or not paper.get("oa_year"):
        return []
    filters = ",".join([
        "publication_year:%d" % paper["oa_year"],
        "primary_topic.field.id:%s" % paper["field"],
        "cited_by_count:>4",
        "is_retracted:false",
    ])
    url = ("%s?filter=%s&sample=3&seed=%s&per-page=3&mailto=%s"
           "&select=id,publication_year,cited_by_count" % (
               API, filters, paper["record_id"], MAIL))
    data = get(url)
    return [{"openalex": w["id"].split("/")[-1],
             "for_record": paper["record_id"],
             "oa_year": w.get("publication_year"),
             "lifetime_cites": w.get("cited_by_count"),
             "retraction_year": paper["retraction_year"],
             "paper_year": paper["paper_year"]}
            for w in (data or {}).get("results", [])]


def main():
    os.makedirs(OUT, exist_ok=True)
    pool = population()
    print("population passing the sealed filters:", len(pool))
    random.Random(SEED).shuffle(pool)

    # Resolve more than we need: not every DOI is in OpenAlex.
    sample, cursor = [], 0
    while len(sample) < TARGET and cursor < len(pool):
        batch = pool[cursor:cursor + 50]
        cursor += 50
        found = resolve(batch)
        for row in batch:
            hit = found.get(row["doi"])
            if hit and hit["oa_year"]:
                row.update(hit)
                sample.append(row)
                if len(sample) == TARGET:
                    break
        print("  resolved %d/%d (asked %d)" % (len(sample), TARGET, cursor),
              end="\r", flush=True)
    print("\nsample:", len(sample), "of", TARGET)
    if len(sample) < 200:
        print("UNDERPOWERED: the rule says no verdict below 200.")

    with ThreadPoolExecutor(max_workers=3) as ex:
        ctl_lists = list(ex.map(controls_for, sample))
    controls = [c for lst in ctl_lists for c in lst]
    print("controls drawn:", len(controls))

    works = ([(p["openalex"], "retracted") for p in sample]
             + [(c["openalex"], "control") for c in controls])
    with ThreadPoolExecutor(max_workers=3) as ex:
        hists = list(ex.map(lambda w: histogram(w[0]), works))

    with open(os.path.join(OUT, "cites.csv"), "w", newline="") as fh:
        w = csv.writer(fh)
        w.writerow(["openalex", "arm", "cite_year", "count"])
        missing = 0
        for (wid, arm), hist in zip(works, hists):
            if hist is None:
                missing += 1
                continue
            for year, count in sorted(hist.items()):
                w.writerow([wid, arm, year, count])
    print("histograms missing:", missing)

    with open(os.path.join(OUT, "sample.csv"), "w", newline="") as fh:
        w = csv.DictWriter(fh, fieldnames=list(sample[0].keys()))
        w.writeheader()
        w.writerows(sample)
    with open(os.path.join(OUT, "controls.csv"), "w", newline="") as fh:
        w = csv.DictWriter(fh, fieldnames=list(controls[0].keys()))
        w.writeheader()
        w.writerows(controls)
    print("written to", OUT)


if __name__ == "__main__":
    main()
