"""
Matched reanalysis of time-to-retraction dispersion.

Referee objection addressed: the machine-attributed cohort is concentrated in recent
retraction years and in a few publishers, so a smaller mean absolute deviation (MAD) may be a
composition artifact. Recent articles cannot yet have long publication-to-retraction intervals.

Stratum definition
------------------
PRIMARY: exact cross of retraction year x publication year x EXACT normalised publisher.
SENSITIVITY 1: the same, but publisher collapsed to the six largest in the cohort plus a pooled
  "Other". This is reported only for comparison, because pooling into "Other" permits a machine
  record from one small publisher to be matched against a comparison record from a different
  small publisher.
SENSITIVITY 2: exact publisher, with Hindawi removed from BOTH cohorts (an earlier
  implementation removed it from the machine cohort only, which is not a valid exclusion test).

Analyses per configuration
--------------------------
  A. Stratum-specific dispersion for strata with >= MIN_PER_ARM records in each arm.
  B. Pooled within-stratum MAD difference, weighted by the smaller arm, tested by permuting the
     cohort label ONLY within strata so the null preserves composition exactly.
  C. A 1:1 sample matched exactly within strata without replacement, compared on MAD and IQR,
     with a within-pair permutation test.

Deterministic given SEED. No network access.
"""

# _REPRO_PATHS_ : resolve paths relative to the repository root, overridable via
# the REPRO_ROOT environment variable. No absolute paths are baked in.
import os as _os
_HERE = _os.path.dirname(_os.path.abspath(__file__))
_ROOT = _os.environ.get("REPRO_ROOT") or _os.path.dirname(_HERE)
def _p(*parts):
    return _os.path.join(_ROOT, *parts)
import csv
import re
import sys
import collections
from datetime import datetime

import numpy as np

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_PERM = 10000
MIN_PER_ARM = 5
REASON_PAT = ("computer-aided content", "computer-generated content")

CSV_PATH = sys.argv[1] if len(sys.argv) > 1 else _p("data", "retraction_watch_gitlab.csv")
OUT_DIR = sys.argv[2] if len(sys.argv) > 2 else _p("runs", "ai-retractions-v3", "tables")


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


def load():
    rows = []
    with open(CSV_PATH, 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 "").lower()
            od = parse_date(row.get("OriginalPaperDate"))
            rd = parse_date(row.get("RetractionDate"))
            if od is None or rd is None:
                continue
            lag = (rd - od).days
            if lag < 0:
                continue
            rows.append({
                "machine": 1 if any(p in reason for p in REASON_PAT) else 0,
                "lag": lag, "ry": rd.year, "py": od.year,
                "publisher": (row.get("Publisher") or "").strip(),
            })
    return rows


def mad(a):
    a = np.asarray(a, dtype=float)
    return float(np.mean(np.abs(a - np.median(a))))


def iqr(a):
    a = np.asarray(a, dtype=float)
    return float(np.percentile(a, 75) - np.percentile(a, 25))


def analyse(rows, stratkey, label, rng_seed, n_seeds=1):
    """Run analyses A, B and C for one stratum definition. Returns (summary, stratum_rows).

    n_seeds > 1 repeats the 1:1 matched draw under different seeds and reports the spread of
    the resulting MAD ratio, so the headline ratio cannot rest on one arbitrary control draw.
    """
    machine = [r for r in rows if r["machine"]]
    control = [r for r in rows if not r["machine"]]
    if not machine or not control:
        return None, []

    m_all = [r["lag"] for r in machine]
    c_all = [r["lag"] for r in control]
    unadj = mad(m_all) / mad(c_all)

    by = collections.defaultdict(lambda: {"m": [], "c": []})
    for r in rows:
        by[(r["ry"], r["py"], r[stratkey])]["m" if r["machine"] else "c"].append(r["lag"])
    inf = {k: v for k, v in by.items()
           if len(v["m"]) >= MIN_PER_ARM and len(v["c"]) >= MIN_PER_ARM}
    if not inf:
        return None, []

    # ---- A. stratum-specific
    srows = []
    for (ry, py, pg), v in sorted(inf.items(), key=lambda kv: str(kv[0])):
        mm, cc = mad(v["m"]), mad(v["c"])
        srows.append({
            "Configuration": label, "RetractionYear": ry, "PublicationYear": py,
            "Publisher": pg, "NMachine": len(v["m"]), "NComparison": len(v["c"]),
            "MedianMachine": round(float(np.median(v["m"])), 1),
            "MedianComparison": round(float(np.median(v["c"])), 1),
            "MadMachine": round(mm, 1), "MadComparison": round(cc, 1),
            "MadRatio": round(mm / cc, 4) if cc > 0 else "",
            "IqrMachine": round(iqr(v["m"]), 1),
            "IqrComparison": round(iqr(v["c"]), 1),
            "IqrRatio": round(iqr(v["m"]) / iqr(v["c"]), 4) if iqr(v["c"]) > 0 else "",
        })
    ratios = [r["MadRatio"] for r in srows if r["MadRatio"] != ""]
    iratios = [r["IqrRatio"] for r in srows if r["IqrRatio"] != ""]
    med_ratio = float(np.median(ratios))          # true median, not the upper-middle element
    n_below = sum(1 for x in ratios if x < 1)
    n_ibelow = sum(1 for x in iratios if x < 1)

    # ---- B. pooled within-stratum + within-stratum permutation
    rng = np.random.default_rng(rng_seed)
    keys = sorted(inf.keys(), key=lambda k: str(k))
    w = np.array([min(len(inf[k]["m"]), len(inf[k]["c"])) for k in keys], float)
    w /= w.sum()
    obs = float(np.sum(w * np.array([mad(inf[k]["m"]) - mad(inf[k]["c"]) for k in keys])))
    stacked = {k: (np.array(inf[k]["m"] + inf[k]["c"], float), len(inf[k]["m"])) for k in keys}
    perm = np.empty(N_PERM)
    for i in range(N_PERM):
        tot = 0.0
        for j, k in enumerate(keys):
            allv, nm = stacked[k]
            idx = rng.permutation(allv.size)
            a, b = allv[idx[:nm]], allv[idx[nm:]]
            tot += w[j] * (np.mean(np.abs(a - np.median(a)))
                           - np.mean(np.abs(b - np.median(b))))
        perm[i] = tot
    p_perm = float((np.sum(np.abs(perm) >= abs(obs)) + 1) / (N_PERM + 1))
    nlo, nhi = np.percentile(perm, [2.5, 97.5])

    # ---- C. 1:1 exact matched without replacement, repeated over seeds
    def one_draw(seed):
        r2 = np.random.default_rng(seed)
        pool = collections.defaultdict(list)
        for r in control:
            pool[(r["ry"], r["py"], r[stratkey])].append(r["lag"])
        for k in pool:
            r2.shuffle(pool[k])
        mm, cc, un = [], [], 0
        for r in machine:
            p = pool.get((r["ry"], r["py"], r[stratkey]))
            if p:
                cc.append(p.pop()); mm.append(r["lag"])
            else:
                un += 1
        return mm, cc, un, r2

    mm_l, cc_l, unmatched, rng2 = one_draw(rng_seed)

    seed_ratios, seed_diffs, seed_pairs = [], [], []
    if n_seeds > 1:
        for s in range(n_seeds):
            a_, b_, _, _ = one_draw(rng_seed + 5000 + s)
            if a_ and mad(b_) > 0:
                seed_ratios.append(mad(a_) / mad(b_))
                seed_diffs.append(mad(a_) - mad(b_))
                seed_pairs.append(len(a_))

    if mm_l:
        d_obs = mad(mm_l) - mad(cc_l)
        a = np.array(mm_l, float); b = np.array(cc_l, float)
        cnt = 0
        for _ in range(N_PERM):
            sw = rng2.random(a.size) < 0.5
            aa = np.where(sw, b, a); bb = np.where(sw, a, b)
            d = (np.mean(np.abs(aa - np.median(aa)))
                 - np.mean(np.abs(bb - np.median(bb))))
            if abs(d) >= abs(d_obs):
                cnt += 1
        p_pair = (cnt + 1) / (N_PERM + 1)
    else:
        d_obs = p_pair = float("nan")

    s = {
        "Configuration": label,
        "MachineRecords": len(machine), "ComparisonRecords": len(control),
        "UnadjustedMadMachine": round(mad(m_all), 1),
        "UnadjustedMadComparison": round(mad(c_all), 1),
        "UnadjustedMadRatio": round(unadj, 4),
        "StrataTotal": len(by), "StrataInformative": len(inf),
        "MachineInInformative": sum(len(v["m"]) for v in inf.values()),
        "ComparisonInInformative": sum(len(v["c"]) for v in inf.values()),
        "StratumMadRatioMedian": round(med_ratio, 4),
        "StrataMadRatioBelow1": f"{n_below}/{len(ratios)}",
        "StrataIqrRatioBelow1": f"{n_ibelow}/{len(iratios)}",
        "PooledWithinStratumMadDiff": round(obs, 2),
        "PermutationP": "< 1e-4" if p_perm < 1e-4 else round(p_perm, 4),
        "PermNullCiLo": round(float(nlo), 2), "PermNullCiHi": round(float(nhi), 2),
        "MatchedPairs": len(mm_l), "MatchedUnmatched": unmatched,
        "MatchedMadMachine": round(mad(mm_l), 1) if mm_l else "",
        "MatchedMadControl": round(mad(cc_l), 1) if mm_l else "",
        "MatchedMadRatio": round(mad(mm_l) / mad(cc_l), 4) if mm_l else "",
        "MatchedIqrMachine": round(iqr(mm_l), 1) if mm_l else "",
        "MatchedIqrControl": round(iqr(cc_l), 1) if mm_l else "",
        "MatchedIqrRatio": round(iqr(mm_l) / iqr(cc_l), 4) if mm_l else "",
        "MatchedPairedMadDiff": round(d_obs, 2) if mm_l else "",
        "MatchedPairedP": ("< 1e-4" if p_pair < 1e-4 else round(p_pair, 4)) if mm_l else "",
        "SeedDraws": len(seed_ratios),
        "SeedRatioMean": round(float(np.mean(seed_ratios)), 4) if seed_ratios else "",
        "SeedRatioSd": (round(float(np.std(seed_ratios, ddof=1)), 4)
                        if len(seed_ratios) > 1 else ""),
        "SeedRatioMin": round(float(np.min(seed_ratios)), 4) if seed_ratios else "",
        "SeedRatioMax": round(float(np.max(seed_ratios)), 4) if seed_ratios else "",
        "SeedPairsMin": min(seed_pairs) if seed_pairs else "",
        "SeedPairsMax": max(seed_pairs) if seed_pairs else "",
    }
    return s, srows


def main():
    rows = load()
    print(f"loaded {len(rows)} retractions with usable dates "
          f"({sum(r['machine'] for r in rows)} machine)")

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

    def norm_pub(p):
        return re.sub(r"\s+", " ", (p or "").strip().lower())

    for r in rows:
        r["pg_exact"] = norm_pub(r["publisher"])
        r["pg_group"] = r["publisher"] if r["publisher"] in top else "Other"
    print(f"distinct normalised publishers: {len({r['pg_exact'] for r in rows})}")
    print(f"publisher groups (sensitivity only): {top}\n")

    hind = [r for r in rows if "hindawi" not in r["pg_exact"]]
    print(f"Hindawi removed from BOTH cohorts: {len(rows)} -> {len(hind)} records "
          f"({sum(r['machine'] for r in hind)} machine)\n")

    configs = [
        (rows, "pg_exact", "PRIMARY: exact publisher"),
        (rows, "pg_group", "SENSITIVITY: publisher group (six largest plus Other)"),
        (hind, "pg_exact", "SENSITIVITY: exact publisher, Hindawi removed from both cohorts"),
    ]

    summaries, all_strata = [], []
    for sub, key, label in configs:
        s, sr = analyse(sub, key, label, SEED, n_seeds=20)
        if s is None:
            print(f"[{label}] no informative strata"); continue
        summaries.append(s); all_strata.extend(sr)
        print(f"[{label}]")
        print(f"   unadjusted MAD ratio        {s['UnadjustedMadRatio']}")
        print(f"   informative strata          {s['StrataInformative']} of {s['StrataTotal']}")
        print(f"   stratum MAD ratio median    {s['StratumMadRatioMedian']} "
              f"({s['StrataMadRatioBelow1']} below 1.0)")
        print(f"   IQR ratios below 1.0        {s['StrataIqrRatioBelow1']}")
        print(f"   pooled within-stratum diff  {s['PooledWithinStratumMadDiff']} days "
              f"(perm p {s['PermutationP']}, null {s['PermNullCiLo']} to {s['PermNullCiHi']})")
        print(f"   matched pairs               {s['MatchedPairs']} "
              f"(unmatched {s['MatchedUnmatched']})")
        print(f"   matched MAD                 {s['MatchedMadMachine']} vs "
              f"{s['MatchedMadControl']}  ratio {s['MatchedMadRatio']}")
        print(f"   matched IQR                 {s['MatchedIqrMachine']} vs "
              f"{s['MatchedIqrControl']}  ratio {s['MatchedIqrRatio']}")
        print(f"   matched paired MAD diff     {s['MatchedPairedMadDiff']} days "
              f"(p {s['MatchedPairedP']})")
        if s["SeedDraws"]:
            print(f"   ratio over {s['SeedDraws']} seeds       "
                  f"mean {s['SeedRatioMean']}  sd {s['SeedRatioSd']}  "
                  f"range {s['SeedRatioMin']} to {s['SeedRatioMax']}  "
                  f"(pairs {s['SeedPairsMin']}-{s['SeedPairsMax']})")
        print()

    with open(f"{OUT_DIR}/t45_stratum_dispersion.csv", "w", newline="",
              encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=list(all_strata[0].keys()),
                           quoting=csv.QUOTE_ALL)
        w.writeheader(); w.writerows(all_strata)
    with open(f"{OUT_DIR}/t46_matched_dispersion_summary.csv", "w", newline="",
              encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=list(summaries[0].keys()),
                           quoting=csv.QUOTE_ALL)
        w.writeheader(); w.writerows(summaries)
    print(f"wrote t45_stratum_dispersion.csv ({len(all_strata)} strata rows)")
    print(f"wrote t46_matched_dispersion_summary.csv ({len(summaries)} configurations)")


if __name__ == "__main__":
    main()
