"""Compute what `rule.retraction-citations-fast` asks for, in its order.

The seal (sha256 00df5043aab13edb6d30aeb4c7f97e5b0362c98bc59851f532ae4af9df69efbc,
2026-09-20T19:52Z) was written before a byte of this population was fetched and
long before this file existed. Everything below is a transcription of it; where
I had a choice the rule did not make, the choice is named in a comment.

Three controls run FIRST and their verdicts are printed before the estimate.
If any fails, the script stops and issues no verdict, as the rule requires.

Writes results.md beside this file. Reads only ~/data/retraction-citations-fast/.
"""

import collections, csv, os, random, statistics

DATA = os.path.expanduser("~/data/retraction-citations-fast")
HERE = os.path.dirname(os.path.abspath(__file__))
BUCKETS = [(1, 2), (3, 5), (6, 10), (11, 20), (21, 50), (51, 200), (201, 10 ** 9)]
BOOTSTRAPS = 2000
SHUFFLES = 200
SEED = 20260921

RR_SLOW = 0.4345          # the published slow-retraction estimate, held fixed
SAME_LO, SAME_HI = 0.85, 1.18
SLOW_BAND = (SAME_LO * RR_SLOW, SAME_HI * RR_SLOW)   # 0.369 .. 0.513


def bucket(pre):
    for lo, hi in BUCKETS:
        if lo <= pre <= hi:
            return "%d-%d" % (lo, hi) if hi < 10 ** 9 else "201+"
    return None


def load():
    """Both arms, with the windows the rule defines.

    P and R are taken from `paper_year` and `retraction_year`, which both files
    carry and which a control row inherits from its retracted paper — that is
    the rule's mechanism for making a pair's windows identical in length. The
    `oa_year`/`cr_year` columns are NOT used for the windows; `cr_year` is
    Crossref's idea of the control journal's volume year and differs from the
    inherited P often enough to matter.
    """
    hist = collections.defaultdict(dict)
    with open(os.path.join(DATA, "cites.csv")) as fh:
        for row in csv.DictReader(fh):
            year = int(row["cite_year"])
            hist[row["doi"]].setdefault("self", 0)
            if year == 0:          # the marker fetch.py writes for "no citations"
                continue
            hist[row["doi"]][year] = int(row["count"])
            hist[row["doi"]]["self"] += int(row["self_citations"])

    papers = []
    with open(os.path.join(DATA, "sample.csv")) as fh:
        for row in csv.DictReader(fh):
            papers.append(_windows(row, hist, "retracted"))
    with open(os.path.join(DATA, "controls.csv")) as fh:
        for row in csv.DictReader(fh):
            papers.append(_windows(row, hist, "control"))
    return [p for p in papers if p is not None]


def _windows(row, hist, arm):
    wid = row["doi"]
    if wid not in hist:
        return None
    r = int(row["retraction_year"])
    p = int(row["paper_year"])
    h = hist[wid]
    pre_years = list(range(p + 1, r))          # P+1 .. R-1, one year or two
    post_years = [r + 1, r + 2, r + 3]         # year R dropped from both
    return {"id": wid, "arm": arm,
            "pre": sum(h.get(y, 0) for y in pre_years),
            "post": sum(h.get(y, 0) for y in post_years),
            "self": h.get("self", 0),
            "P": p, "R": r, "pre_len": len(pre_years),
            "record": row.get("record_id") or row.get("for_record")}


def strata(papers):
    """(publication year, pre-window length, pre bucket) -> {arm: [papers]}.

    The triple is the rule's, and the middle term is the whole reason there is
    a second study: a pre-window of one year and one of two years are not
    comparable quantities, so no stratum is allowed to hold both.
    """
    out = collections.defaultdict(lambda: {"retracted": [], "control": []})
    for p in papers:
        if p["pre"] <= 0:
            continue
        out[(p["P"], p["pre_len"], bucket(p["pre"]))][p["arm"]].append(p)
    return out


def rr_mh(cells, arm_a="retracted", arm_b="control"):
    """RR = sum_s (A_s M_s / T_s) / sum_s (B_s N_s / T_s), written in the seal.

    Nothing is divided inside a stratum. Returns (RR, strata used, retracted
    pre-window mass covered, total retracted pre-window mass).
    """
    num = den = 0.0
    used = 0
    covered = total = 0
    for cell in cells.values():
        a, b = cell[arm_a], cell[arm_b]
        n_a = sum(p["pre"] for p in a)
        n_b = sum(p["pre"] for p in b)
        total += n_a
        if not a or not b or n_a == 0 or n_b == 0:
            continue
        t = n_a + n_b
        num += sum(p["post"] for p in a) * n_b / t
        den += sum(p["post"] for p in b) * n_a / t
        used += 1
        covered += n_a
    return (num / den if den else float("nan")), used, covered, total


def bootstrap(cells, rng, n=BOOTSTRAPS):
    """Resample papers within stratum and arm, as the rule says."""
    out = []
    keys = list(cells)
    for _ in range(n):
        draw = {}
        for k in keys:
            cell = cells[k]
            draw[k] = {arm: [rng.choice(cell[arm]) for _ in cell[arm]]
                       if cell[arm] else [] for arm in ("retracted", "control")}
        value, _, _, _ = rr_mh(draw)
        if value == value:
            out.append(value)
    out.sort()
    return out[int(0.025 * len(out))], out[int(0.975 * len(out))]


def null_a(papers, rng):
    """Composition-preserving: labels shuffled WITHIN each stratum."""
    cells = strata(papers)
    values = []
    for _ in range(SHUFFLES):
        draw = {}
        for k, cell in cells.items():
            pooled = cell["retracted"] + cell["control"]
            rng.shuffle(pooled)
            n = len(cell["retracted"])
            draw[k] = {"retracted": pooled[:n], "control": pooled[n:]}
        value, _, _, _ = rr_mh(draw)
        if value == value:
            values.append(value)
    return statistics.mean(values), min(values), max(values)


def null_b(papers, rng):
    """Composition-breaking: labels shuffled across the whole sample, strata
    then rebuilt from each paper's own (P, pre-window length, bucket).

    Vex of the Drift's point: the permutation scheme is the hypothesis. A runs
    against "within a stratum the arms are exchangeable"; B against "the arm
    label carries nothing at all". An estimator that leans on how the arms are
    spread across strata passes one and fails the other.
    """
    n_ret = sum(1 for p in papers if p["arm"] == "retracted")
    pool = list(papers)
    values = []
    for _ in range(SHUFFLES):
        rng.shuffle(pool)
        fake = [dict(p, arm="retracted" if i < n_ret else "control")
                for i, p in enumerate(pool)]
        value, _, _, _ = rr_mh(strata(fake))
        if value == value:
            values.append(value)
    return statistics.mean(values), min(values), max(values)


def positive_control(papers, rng):
    """Half the control arm, post-window thinned by 50 %, must come back ~0.5."""
    ctl = [p for p in papers if p["arm"] == "control"]
    rng.shuffle(ctl)
    half = len(ctl) // 2
    fake = []
    for i, p in enumerate(ctl):
        q = dict(p)
        if i < half:
            q["arm"] = "retracted"
            q["post"] = int(q["post"] * 0.5)
        else:
            q["arm"] = "control"
        fake.append(q)
    cells = strata(fake)
    value, used, _, _ = rr_mh(cells)
    lo, hi = bootstrap(cells, rng, n=500)
    return value, lo, hi, used


def main():
    rng = random.Random(SEED)
    papers = load()
    ret = [p for p in papers if p["arm"] == "retracted"]
    ctl = [p for p in papers if p["arm"] == "control"]
    lines = []

    def say(text=""):
        print(text)
        lines.append(text)

    say("# Does a fast retraction bite as hard as a slow one? — results")
    say()
    say("Computed by `analyse.py` from `~/data/retraction-citations-fast/`. "
        "The rule was sealed as `rule.retraction-citations-fast` at "
        "2026-09-20T19:52Z, before any of this population was fetched and "
        "before this file existed.")
    say()

    say("## The sample")
    say()
    drop_ret = sum(1 for p in ret if p["pre"] <= 0)
    drop_ctl = sum(1 for p in ctl if p["pre"] <= 0)
    say("- retracted papers resolved: **%d** (the rule's floor for a verdict "
        "is 150)" % len(ret))
    say("- control papers (same journal, same publication year): **%d**"
        % len(ctl))
    say("- **retracted papers with an EMPTY pre-window, dropped by the rule: "
        "%d of %d (%.1f %%)**" % (drop_ret, len(ret),
                                  100.0 * drop_ret / max(1, len(ret))))
    say("- control papers likewise dropped: **%d of %d (%.1f %%)**"
        % (drop_ctl, len(ctl), 100.0 * drop_ctl / max(1, len(ctl))))
    say()
    say("The rule names that drop as the bias this population carries and asks "
        "for it prominently: a paper pulled two or three years after "
        "publication has one or two years in which to be cited before the "
        "notice lands, so the papers that vanish here are the least-cited "
        "ones — plausibly the ones a retraction would have silenced most "
        "completely. The direction of that bias is toward RR = 1, against the "
        "result I expected.")
    say()
    lens = collections.Counter(p["pre_len"] for p in ret if p["pre"] > 0)
    say("- retracted papers by pre-window length: "
        + ", ".join("**%d year%s: %d**" % (k, "" if k == 1 else "s", v)
                    for k, v in sorted(lens.items())))
    say()
    if len(ret) < 150:
        say("**Fewer than 150 retracted papers resolved. The rule reports the "
            "study as underpowered and issues no verdict.**")
        _write(lines)
        return

    say("## The controls on the detector")
    say()
    a_mean, a_lo, a_hi = null_a(papers, rng)
    a_ok = 0.90 <= a_mean <= 1.11
    say("**Null control A, composition-preserving** — arm labels shuffled "
        "inside each stratum, %d draws: mean RR **%.4f** (range %.3f to "
        "%.3f), band 0.90 to 1.11: **%s**."
        % (SHUFFLES, a_mean, a_lo, a_hi, "PASS" if a_ok else "FAIL"))
    say()
    b_mean, b_lo, b_hi = null_b(papers, rng)
    b_ok = 0.90 <= b_mean <= 1.11
    say("**Null control B, composition-breaking** — labels shuffled "
        "across the whole sample and the strata rebuilt, %d draws: mean RR "
        "**%.4f** (range %.3f to %.3f), band 0.90 to 1.11: **%s**."
        % (SHUFFLES, b_mean, b_lo, b_hi, "PASS" if b_ok else "FAIL"))
    say()
    pos, plo, phi, pused = positive_control(papers, rng)
    pos_ok = 0.42 <= pos <= 0.58
    say("**Positive control** — half the control arm thinned by 50 %%, "
        "over %d strata: RR **%.4f** (95 %% CI %.3f to %.3f), band 0.42 to "
        "0.58: **%s**." % (pused, pos, plo, phi, "PASS" if pos_ok else "FAIL"))
    say()
    floor = phi - plo
    say("**Floor.** The positive control's interval is %.3f wide, so this "
        "sample cannot distinguish an effect smaller than about %.0f %% from "
        "none at all." % (floor, 100 * floor / 2))
    say()
    if not (a_ok and b_ok and pos_ok):
        say("**A control failed. No verdict is issued, on either the primary "
            "or the secondary claim.**")
        _write(lines)
        return

    say("## The estimate")
    say()
    cells = strata(papers)
    value, used, covered, total = rr_mh(cells)
    lo, hi = bootstrap(cells, rng)
    say("Strata used: **%d** of %d, holding **%.1f %%** of the retracted "
        "pre-window citation mass; the rest sat in strata with no control arm "
        "and was dropped." % (used, len(cells),
                              100.0 * covered / total if total else 0))
    say()
    say("**RR = %.4f** (95 %% CI %.3f to %.3f)." % (value, lo, hi))
    say()

    say("### Primary verdict")
    say()
    if value >= 0.85 and lo > 0.85:
        verdict = ("SURVIVED", "A fast retraction does not measurably change "
                               "the rate at which a paper is cited.")
    elif value <= 0.50 and hi < 0.50:
        verdict = ("KILLED", "A fast retraction at least halves the citation "
                             "rate, as a slow one does.")
    else:
        verdict = ("INCONCLUSIVE",
                   "Between the sealed bars: the citation rate falls without "
                   "the fall clearing either threshold.")
    say("**%s.** %s" % verdict)
    say()

    say("### Secondary verdict, declaredly the weaker of the two")
    say()
    d = value / RR_SLOW
    gap = abs(value - RR_SLOW)
    say("RR_slow is held FIXED at the published **%.4f** and is not "
        "re-estimated, so the uncertainty below is the fast study's alone. "
        "D = RR_fast / RR_slow = **%.4f**." % (RR_SLOW, d))
    say()
    if floor > gap:
        second = ("UNDECIDED", "The floor (%.3f) is wider than the gap between "
                               "the two estimates (%.3f), so this sample could "
                               "not have told the two populations apart "
                               "whatever it found." % (floor, gap))
    elif not (lo <= RR_SLOW <= hi):
        second = ("DIFFERENT", "The fast study's interval (%.3f to %.3f) "
                               "excludes the slow study's point estimate."
                               % (lo, hi))
    elif (SAME_LO <= d <= SAME_HI
          and SLOW_BAND[0] <= lo and hi <= SLOW_BAND[1]):
        second = ("SAME", "D is inside [0.85, 1.18] and the whole interval "
                          "lies inside [%.3f, %.3f]." % SLOW_BAND)
    else:
        second = ("UNDECIDED", "Neither the exclusion nor the equivalence "
                               "condition is met.")
    say("**%s.** %s" % second)
    say()

    say("## The plain counts, which carry no verdict")
    say()
    cited = [p for p in ret if p["post"] > 0]
    say("- fast-retracted papers collecting at least one citation in the three "
        "years after the retraction year: **%d of %d (%.1f %%)**"
        % (len(cited), len(ret), 100.0 * len(cited) / len(ret)))
    say("- median post-window citations, retracted: **%d**; controls: **%d**"
        % (statistics.median([p["post"] for p in ret]),
           statistics.median([p["post"] for p in ctl])))
    say("- total citations to these %d fast-retracted papers in the three "
        "years after their retraction: **%d**"
        % (len(ret), sum(p["post"] for p in ret)))
    say("- of all citations to the retracted papers, OpenCitations flags "
        "**%d** as author self-citations (%.1f %%)."
        % (sum(p["self"] for p in ret),
           100.0 * sum(p["self"] for p in ret)
           / max(1, sum(p["self"] for p in ret)
                 + sum(p["pre"] + p["post"] for p in ret))))
    say()

    say("## Per stratum")
    say()
    say("| publication year | pre-window | pre-window citations | retracted | "
        "controls | post/pre retracted | post/pre control |")
    say("|---|---|---|---|---|---|---|")
    for (year, plen, b), cell in sorted(cells.items()):
        a, c = cell["retracted"], cell["control"]
        if not a or not c:
            continue
        pa, pc = sum(p["pre"] for p in a), sum(p["pre"] for p in c)
        if pa == 0 or pc == 0:
            continue
        say("| %d | %d year%s | %s | %d | %d | %.3f | %.3f |"
            % (year, plen, "" if plen == 1 else "s", b, len(a), len(c),
               sum(p["post"] for p in a) / pa, sum(p["post"] for p in c) / pc))
    _write(lines)


def _write(lines):
    with open(os.path.join(HERE, "results.md"), "w") as fh:
        fh.write("\n".join(lines) + "\n")
    print("\nwritten:", os.path.join(HERE, "results.md"))


if __name__ == "__main__":
    main()
