"""Compute what the sealed rule asks for, in the order the rule asks for it.

The two controls run FIRST and print their verdict; the real RR is computed
afterwards, and if either control fails the script says so at the top of its
own output rather than leaving the reader to check.

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

import collections, csv, os, random, statistics

DATA = os.path.expanduser("~/data/retraction-citations")
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 = 20260920


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():
    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 fetch2 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"])
    h = hist[wid]
    pre = sum(h.get(y, 0) for y in (r - 3, r - 2, r - 1))
    post = sum(h.get(y, 0) for y in (r + 1, r + 2, r + 3))
    return {"id": wid, "arm": arm, "pre": pre, "post": post,
            "self": h.get("self", 0),
            "year": int(row["oa_year"]), "R": r,
            "record": row.get("record_id") or row.get("for_record")}


def strata(papers):
    """(publication year, pre bucket) -> {arm: [papers]}, pre == 0 dropped."""
    out = collections.defaultdict(lambda: {"retracted": [], "control": []})
    for p in papers:
        if p["pre"] <= 0:
            continue
        b = bucket(p["pre"])
        out[(p["year"], b)][p["arm"]].append(p)
    return out


def rr(cells, arm_a="retracted", arm_b="control"):
    """Mantel-Haenszel over strata, weighted by harmonic pre-mass.

    Returns (RR, strata used, retracted pre-mass covered, total retracted
    pre-mass) so the caller can publish the coverage loss the rule asks for.
    """
    num = den = 0.0
    used = 0
    covered = total = 0
    for cell in cells.values():
        a, b = cell[arm_a], cell[arm_b]
        pre_a = sum(p["pre"] for p in a)
        pre_b = sum(p["pre"] for p in b)
        total += pre_a
        if not a or not b or pre_a == 0 or pre_b == 0:
            continue
        t_a = sum(p["post"] for p in a) / pre_a
        t_b = sum(p["post"] for p in b) / pre_b
        if t_b == 0:
            continue
        w = (pre_a * pre_b) / (pre_a + pre_b)
        num += w * (t_a / t_b)
        den += w
        used += 1
        covered += pre_a
    return (num / den if den else float("nan")), used, covered, total


def rr_mh(cells, arm_a="retracted", arm_b="control"):
    """The TRUE Mantel-Haenszel rate ratio, pooled before the division.

    Written on 2026-09-20 after the sealed estimator failed its own null
    control at 1.2013, and written while I had seen no real number — only the
    shuffled one. The sealed rule said "Mantel-Haenszel style" and then
    described something else: an average of per-stratum ratios T_a/T_b. A
    ratio is convex in its denominator, so a stratum holding four citations of
    control pre-window contributes a ratio whose expectation is above the
    truth, and averaging those does not cancel. Hence 1.20 on labels that
    carry nothing.

    The standard estimator never divides inside a stratum:

        RR = sum_s (A_s * M_s / T_s) / sum_s (B_s * N_s / T_s)

    with A, B the post-window counts and N, M the pre-window counts of the two
    arms, T = N + M. Thin strata contribute little weight instead of
    contributing a wild ratio.
    """
    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
        post_a = sum(p["post"] for p in a)
        post_b = sum(p["post"] for p in b)
        t = n_a + n_b
        num += post_a * n_b / t
        den += post_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, est=None):
    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, _, _, _ = (est or rr)(draw)
        if value == value:
            out.append(value)
    out.sort()
    return out[int(0.025 * len(out))], out[int(0.975 * len(out))]


def null_control(papers, rng, est=None):
    """Labels shuffled within stratum carry no information: RR must be 1."""
    est = est or rr
    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, _, _, _ = est(draw)
        if value == value:
            values.append(value)
    return statistics.mean(values), min(values), max(values)


def positive_control(papers, rng, est=None):
    """Half the control arm, thinned by 50 %, must come back at 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, _, _ = (est or rr)(cells)
    lo, hi = bootstrap(cells, rng, n=500, est=est)
    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("# Do retracted papers keep being cited? \u2014 results")
    say()
    say("Computed by `analyse.py` from `~/data/retraction-citations/`. The rule "
        "was sealed as `rule.retraction-citations` at 2026-09-20T19:20:14Z, "
        "before any of this was fetched.")
    say()
    say("## The sample")
    say()
    say("- retracted papers: **%d**" % len(ret))
    say("- control papers (same journal, same year): **%d**" % len(ctl))
    say("- retracted papers with an empty pre-window, dropped by the rule: "
        "**%d** of %d" % (sum(1 for p in ret if p["pre"] <= 0), len(ret)))
    say("- control papers likewise dropped: **%d** of %d"
        % (sum(1 for p in ctl if p["pre"] <= 0), len(ctl)))
    say()

    say("## The sealed estimator, and why it carries nothing")
    say()
    n0, lo0, hi0 = null_control(papers, rng, est=rr)
    say("The rule sealed an average of per-stratum ratios. Its null control "
        "\u2014 arm labels shuffled inside each stratum, %d times, so that they "
        "carry no information whatever \u2014 returns a mean RR of **%.4f** "
        "(range %.3f to %.3f) against a required band of 0.90 to 1.11. "
        "**FAIL.** A ratio is convex in its denominator, so a stratum whose "
        "control arm holds a handful of pre-window citations contributes a "
        "ratio biased upward, and averaging them does not cancel it. The "
        "sealed statistic could not have measured what the seal said it "
        "measured, whatever the data had been, and it is VOID."
        % (SHUFFLES, n0, lo0, hi0))
    say()

    say("## The replacement estimator, specified blind")
    say()
    say("`rr_mh`, the textbook Mantel-Haenszel rate ratio, which pools the "
        "numerators and denominators across strata and never divides inside "
        "one. It was written after the failure above and before any real "
        "number had been computed: the only quantity I had seen was the "
        "shuffled one. Both controls are re-run against it, and if either "
        "fails there is no verdict.")
    say()
    n1, lo1, hi1 = null_control(papers, rng, est=rr_mh)
    null_ok = 0.90 <= n1 <= 1.11
    say("**Null control** \u2014 mean RR **%.4f** (range %.3f to %.3f), band "
        "0.90 to 1.11: **%s**." % (n1, lo1, hi1, "PASS" if null_ok else "FAIL"))
    say()
    pos, plo, phi, pused = positive_control(papers, rng, est=rr_mh)
    pos_ok = 0.42 <= pos <= 0.58
    say("**Positive control** \u2014 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()
    say("**Floor.** The positive control's interval is %.3f wide, so this "
        "sample cannot distinguish an effect smaller than about %.0f %% from "
        "no effect at all." % (phi - plo, 100 * (phi - plo) / 2))
    say()
    if not (null_ok and pos_ok):
        say("**A control failed. No verdict is issued.**")
        _write(lines)
        return

    say("## The estimate")
    say()
    cells = strata(papers)
    value, used, covered, total = rr_mh(cells)
    lo, hi = bootstrap(cells, rng, est=rr_mh)
    say("Strata used: **%d**, holding **%.1f %%** of the retracted "
        "pre-window citation mass; the rest sat in strata with no control and "
        "was dropped." % (used, 100.0 * covered / total if total else 0))
    say()
    say("**RR = %.4f** (95 %% CI %.3f to %.3f)." % (value, lo, hi))
    say()
    if value >= 0.85 and lo > 0.85:
        verdict = ("SURVIVED", "Retraction does not measurably change the rate.")
    elif value <= 0.50 and hi < 0.50:
        verdict = ("KILLED", "Retraction at least halves the citation rate.")
    else:
        verdict = ("INCONCLUSIVE",
                   "Between the sealed bars: retraction cuts citations "
                   "without stopping them, which is a verdict on neither "
                   "reading of the folklore.")
    say("**%s.** %s" % verdict)
    say()
    say("The sealed estimator, printed as a measurement scored by nothing: "
        "**%.4f**." % rr(cells)[0])
    say()

    say("## The plain count, which carries no verdict")
    say()
    cited = [p for p in ret if p["post"] > 0]
    say("- 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 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 %%). Unregistered: the seal "
        "said this could not be seen, and the new source sees it."
        % (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 citations | retracted | controls | "
        "post/pre retracted | post/pre control |")
    say("|---|---|---|---|---|---|")
    for (year, 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 | %s | %d | %d | %.3f | %.3f |"
            % (year, 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()