"""Collect the sealed fast-retraction population and its raw citation dates.

The rule (`rule.retraction-citations-fast`, sha256
00df5043aab13edb6d30aeb4c7f97e5b0362c98bc59851f532ae4af9df69efbc, sealed
2026-09-20T19:52Z) was written and sealed before this ran, and this file does
not implement any part of it that involves a choice. It stores **raw** citation
year counts per DOI; the windows, the strata and the estimator all live in the
analysis, which has not been written yet. So a session that runs this has not
thereby decided anything the seal left open.

Population, per the seal: Retraction Watch, `RetractionNature == "Retraction"`,
R in 2018..2021, **P in {R-3, R-2}** (the filter that differs from the finished
study), no reinstatement, unique DOI. 300 drawn at seed 20260921.

Sources are Crossref (ISSN, year, and a true random draw of same-journal
same-year controls) and OpenCitations (every citation with its date). Never
OpenAlex: its keyless budget is 1,000 requests a day and this needs more.
"""


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

MAIL = "vesper@untilnextsession.com"
OUT = os.path.expanduser("~/data/retraction-citations-fast")
SEED = 20260921
TARGET = 300

import importlib.util as _ilu
_spec = _ilu.spec_from_file_location("fetch1", os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
    "retraction-citations", "fetch.py"))
_f1 = _ilu.module_from_spec(_spec)
_spec.loader.exec_module(_f1)
year_of = _f1.year_of
RW = _f1.RW


def population():
    """The sealed fast filters. Identical to the finished study except for the
    one line the seal changes: `P in {R-3, R-2}` instead of `P <= R-4`."""
    import csv as _csv
    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):
            continue
        if p not in (r - 3, r - 2):        # <-- the sealed fast filter
            continue
        doi = (row["OriginalPaperDOI"] or "").strip().lower()
        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

_throttle = threading.Semaphore(4)
_pace, _last = threading.Lock(), [0.0]


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


def get(url, tries=3):
    for attempt in range(tries):
        with _throttle:
            _turn()
            try:
                req = urllib.request.Request(
                    url, headers={"User-Agent": "vesper/1.0 (mailto:%s)" % MAIL})
                with urllib.request.urlopen(req, timeout=45) as r:
                    return json.loads(r.read().decode())
            except urllib.error.HTTPError as exc:
                if exc.code == 404:
                    return None
                if attempt == tries - 1:
                    return None
                time.sleep(2.0 * (attempt + 1))
            except Exception:
                if attempt == tries - 1:
                    return None
                time.sleep(2.0 * (attempt + 1))
    return None


def crossref_record(doi):
    """ISSN and published year for one DOI, or None if Crossref has neither."""
    data = get("https://api.crossref.org/works/%s?mailto=%s"
               % (urllib.parse.quote(doi, safe="/."), MAIL))
    msg = (data or {}).get("message")
    if not msg:
        return None
    issns = msg.get("ISSN") or []
    parts = ((msg.get("published") or {}).get("date-parts")
             or (msg.get("issued") or {}).get("date-parts") or [[None]])
    year = parts[0][0] if parts and parts[0] else None
    return {"issn": issns[0] if issns else None, "cr_year": year,
            "journal": (msg.get("container-title") or [None])[0]}


def controls_for(paper):
    """Three random articles from the same journal in the same year.

    Crossref's `sample` is a true random draw over the filtered set, so this
    conditions on nothing but venue and vintage.
    """
    if not paper.get("issn") or not paper.get("cr_year"):
        return []
    year = int(paper["cr_year"])
    url = ("https://api.crossref.org/works?filter=issn:%s,type:journal-article,"
           "from-pub-date:%d-01-01,until-pub-date:%d-12-31&sample=3"
           "&select=DOI,issued&mailto=%s" % (paper["issn"], year, year, MAIL))
    data = get(url)
    items = ((data or {}).get("message") or {}).get("items") or []
    out = []
    for it in items:
        doi = (it.get("DOI") or "").lower()
        if not doi or doi == paper["doi"]:
            continue
        out.append({"doi": doi, "for_record": paper["record_id"],
                    "oa_year": year, "retraction_year": paper["retraction_year"],
                    "paper_year": paper["paper_year"]})
    return out


def citations(doi):
    """[(citing doi, creation date, author self-citation flag)] or None."""
    data = get("https://opencitations.net/index/api/v1/citations/%s"
               % urllib.parse.quote(doi, safe="/."))
    if data is None:
        return None
    return [(c.get("citing"), c.get("creation") or "", c.get("author_sc") or "")
            for c in data]


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

    sample, cursor = [], 0
    with ThreadPoolExecutor(max_workers=4) as ex:
        while len(sample) < TARGET and cursor < len(pool):
            batch = pool[cursor:cursor + 60]
            cursor += 60
            for row, rec in zip(batch, ex.map(
                    lambda r: crossref_record(r["doi"]), batch)):
                if rec and rec["issn"] and rec["cr_year"]:
                    row.update(rec)
                    row["oa_year"] = row["paper_year"]
                    sample.append(row)
            print("  crossref: %d/%d resolved from %d asked"
                  % (len(sample), TARGET, cursor), flush=True)
    sample = sample[:TARGET]
    print("sample:", len(sample), flush=True)

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

    works = ([(p["doi"], "retracted") for p in sample]
             + [(c["doi"], "control") for c in controls])
    with ThreadPoolExecutor(max_workers=4) as ex:
        cites = list(ex.map(lambda w: citations(w[0]), works))

    missing = 0
    with open(os.path.join(OUT, "cites.csv"), "w", newline="") as fh:
        w = csv.writer(fh)
        w.writerow(["doi", "arm", "cite_year", "count", "self_citations"])
        for (doi, arm), rows in zip(works, cites):
            if rows is None:
                missing += 1
                continue
            per_year, selfs = {}, {}
            for _citing, created, sc in rows:
                year = created[:4]
                if not year.isdigit():
                    continue
                per_year[int(year)] = per_year.get(int(year), 0) + 1
                if sc == "yes":
                    selfs[int(year)] = selfs.get(int(year), 0) + 1
            # A paper with no citations at all writes one zero row, so that a
            # real zero and a failed request stay distinguishable downstream.
            if not per_year:
                w.writerow([doi, arm, 0, 0, 0])
            for year, count in sorted(per_year.items()):
                w.writerow([doi, arm, year, count, selfs.get(year, 0)])
    print("citation lookups that failed:", missing, flush=True)

    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, flush=True)


if __name__ == "__main__":
    main()
