"""
Dependence-aware uncertainty for the Mantel-Haenszel paper-mill co-coding estimate.

The Greenland-Robins interval assumes records are independent. Records sharing a retraction
notice are not, so that interval understates dependence. This script recomputes the
Mantel-Haenszel risk ratio and odds ratio under a nonparametric bootstrap that resamples whole
retraction notices, giving an interval that respects the clustering the manuscript documents.

Strata are retraction year crossed with publisher group (six largest in the cohort plus a
residual), matching the specification in Section 3.5. Records whose retraction DOI is blank or
the literal 'unavailable' are singleton clusters.

Writes t51_mh_cluster_bootstrap.csv.
"""

import csv
import os
import collections
from datetime import datetime

import numpy as np

_HERE = os.path.dirname(os.path.abspath(__file__))
_ROOT = os.environ.get("REPRO_ROOT") or os.path.dirname(_HERE)
RAW = os.path.join(_ROOT, "data", "retraction_watch_gitlab.csv")
OUT = os.path.join(_ROOT, "runs", "ai-retractions-v3", "tables")

csv.field_size_limit(10**9)
# CANONICAL date formats, identical in every script in this package. The AM/PM variant is
# required by exactly one record (ID 18930, a 1756 retraction); omitting it silently drops
# that record and yields 66,012 rather than 66,013 analysed retractions.
_DATE_FORMATS = ("%m/%d/%Y %I:%M:%S %p", "%m/%d/%Y %H:%M:%S",
                 "%m/%d/%Y %H:%M", "%m/%d/%Y", "%Y-%m-%d", "%d/%m/%Y")
SEED = 20260731
N_BOOT = 10000
REASON_PAT = ("computer-aided content", "computer-generated content")


def parse_date(s):
    s = (s or "").strip()
    for fmt in _DATE_FORMATS:
        try:
            return datetime.strptime(s, fmt)
        except ValueError:
            continue
    return None


def mh(tab):
    """tab: list of (a, b, c, d) per stratum. a=exposed+, b=exposed-, c=unexp+, d=unexp-."""
    rr_num = rr_den = or_num = or_den = 0.0
    for a, b, c, d in tab:
        n1, n0 = a + b, c + d
        n = n1 + n0
        if n == 0 or n1 == 0 or n0 == 0:
            continue
        rr_num += a * n0 / n
        rr_den += c * n1 / n
        or_num += a * d / n
        or_den += b * c / n
    rr = rr_num / rr_den if rr_den > 0 else float("nan")
    orr = or_num / or_den if or_den > 0 else float("nan")
    return rr, orr


def main():
    rows = []
    with open(RAW, encoding="utf-8-sig", newline="") as f:
        for row in csv.DictReader(f):
            if (row.get("RetractionNature") or "").strip() != "Retraction":
                continue
            reason = (row.get("Reason") or "")
            low = reason.lower()
            rd = parse_date(row.get("RetractionDate"))
            if rd is None:
                continue
            raw = (row.get("RetractionDOI") or "").strip()
            key = (f"__singleton_{row.get('Record ID')}"
                   if raw == "" or raw.lower() == "unavailable" else raw.lower())
            rows.append({
                "cluster": key,
                "machine": any(p in low for p in REASON_PAT),
                "mill": "paper mill" in low,
                "ry": rd.year,
                "publisher": (row.get("Publisher") or "").strip(),
            })

    top = [p for p, _ in collections.Counter(
        r["publisher"] for r in rows if r["machine"]).most_common(6)]
    for r in rows:
        r["stratum"] = (r["ry"], r["publisher"] if r["publisher"] in top else "Other")

    clusters = collections.defaultdict(list)
    for r in rows:
        clusters[r["cluster"]].append(r)
    keys = sorted(clusters)
    print(f"records {len(rows)}  clusters {len(keys)}  strata "
          f"{len({r['stratum'] for r in rows})}")

    def tabulate(recs):
        agg = collections.defaultdict(lambda: [0, 0, 0, 0])
        for r in recs:
            cell = agg[r["stratum"]]
            if r["machine"]:
                cell[0 if r["mill"] else 1] += 1
            else:
                cell[2 if r["mill"] else 3] += 1
        return list(agg.values())

    obs_rr, obs_or = mh(tabulate(rows))
    contributing = sum(1 for a, b, c, d in tabulate(rows)
                       if (a + b) > 0 and (c + d) > 0)
    print(f"observed MH risk ratio {obs_rr:.4f}   odds ratio {obs_or:.4f}   "
          f"strata with both cohorts {contributing}")

    rng = np.random.default_rng(SEED)
    idx = np.arange(len(keys))
    boot = np.empty((N_BOOT, 2))
    for i in range(N_BOOT):
        pick = rng.choice(idx, size=len(keys), replace=True)
        recs = []
        for j in pick:
            recs.extend(clusters[keys[j]])
        boot[i] = mh(tabulate(recs))

    out = []
    for j, (name, point) in enumerate((("Mantel-Haenszel risk ratio", obs_rr),
                                       ("Mantel-Haenszel odds ratio", obs_or))):
        col = boot[~np.isnan(boot[:, j]), j]
        lo, hi = np.percentile(col, [2.5, 97.5])
        q = [np.percentile(col[k::4], [2.5, 97.5]) for k in range(4)]
        mc_lo = float(np.std([x[0] for x in q], ddof=1)) / 2
        mc_hi = float(np.std([x[1] for x in q], ddof=1)) / 2
        print(f"{name}: {point:.4f}  cluster-robust 95% CI {lo:.4f} to {hi:.4f}  "
              f"(MC SE {mc_lo:.4f}, {mc_hi:.4f})")
        out.append({"Estimate": name, "PointEstimate": round(float(point), 4),
                    "ClusterCiLo": round(float(lo), 4), "ClusterCiHi": round(float(hi), 4),
                    "McSeCiLo": round(mc_lo, 4), "McSeCiHi": round(mc_hi, 4),
                    "Clusters": len(keys), "StrataWithBothCohorts": contributing,
                    "Replicates": N_BOOT,
                    "Note": "Resamples whole retraction notices. Compare with the "
                            "Greenland-Robins interval, which assumes independence. "
                            "StrataWithBothCohorts counts strata containing at least one "
                            "record in each cohort, which is the condition for a stratum to "
                            "enter the Mantel-Haenszel sums; it is NOT the same as the 23 "
                            "analytically contributing strata reported in Table 10 of the "
                            "manuscript, which additionally requires a non-zero paper-mill "
                            "numerator, nor the 8 strata with all four cells non-zero used "
                            "by the Cochran Q test. The three counts answer three different "
                            "questions and are expected to differ."})

    path = os.path.join(OUT, "t51_mh_cluster_bootstrap.csv")
    with open(path, "w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=list(out[0].keys()), quoting=csv.QUOTE_ALL)
        w.writeheader(); w.writerows(out)
    print(f"wrote {path}")


if __name__ == "__main__":
    main()
