#!/usr/bin/env python3
"""fred_two_cues.py — the pre-registered two-cue test on FReD (survives-a-recheck, Stage 0 object).

Asks two questions of the FORRT Replication Database (FReD), pinned at the OSF file of
2026-02-01 (sha256 59545e3f…), and answers them by the rules in
`PREREG — survives-a-recheck — FReD two cues — 2026-10-26.md`:

  H1  Do findings whose original p was below .001 replicate more often than findings
      whose original p sat between .01 and .05?
  H2  Do findings from the largest original samples (top third) replicate more often
      than findings from the smallest (bottom third)?

The outcome is the replicating authors' own verdict (`reported_success`). Each original
paper counts once: a row's weight is 1 / (rows from the same original DOI in the set).
Intervals are a cluster bootstrap over original DOIs.

    python3 fred_two_cues.py            # the real run (after the pre-registration is locked)
    python3 fred_two_cues.py --dummy    # never reads an outcome column; random outcomes

Exit 0 on a completed run; 1 if the data file does not match the pinned hash.
"""
from __future__ import annotations

import hashlib
import json
import os
import sys

import numpy as np
import pandas as pd

HERE = os.path.dirname(os.path.abspath(__file__))
DATA = os.path.join(HERE, "data", "fred-2026-02-01-subset.csv")
DATA_SHA256 = "f3e729378bb46c8cf8896b9d85711b1535852f4a36d05bcdeb6848f972ef630c"
RESULTS = os.path.join(HERE, "fred-two-cues-results.json")
SEED = 20261026
BOOT = 5000
OSC_DOI = "aac4716"  # Open Science Collaboration 2015, Science, doi 10.1126/science.aac4716
OUTCOME_COLS = ["reported_success", "pval_value_r", "pval_type_r"]
MIN_DIFF_PP = 10.0  # a cue "works" only if it moves the rate by at least ten points


def norm_type(t) -> str | None:
    """FReD writes the p-value relation several ways ('l (<)', '<', 'e (=)', '"="', ...)."""
    if t is None or (isinstance(t, float) and np.isnan(t)):
        return None
    s = str(t).strip().strip('"')
    if "<=" in s or s.startswith("le"):
        return "<="
    if ">" in s or s.startswith("g"):
        return ">"
    if "<" in s or s.startswith("l"):
        return "<"
    if "=" in s or s.startswith("e"):
        return "="
    return None


def p_band(t, v) -> str:
    """Original p-value band. Only bands a report can be placed in without guessing.

    S  exact p < .001, or reported as below a value of .001 or less
    M  exact .001 <= p < .01
    W  exact .01 <= p <= .05
    UNRES_M / UNRES_W  reported only as "below .01" / "below .05" (could be a lower band)
    NS original not significant (exact p > .05, or reported as above a value)
    MISSING  no usable p
    """
    kind = norm_type(t)
    if kind is None or v is None or (isinstance(v, float) and np.isnan(v)):
        return "MISSING"
    v = float(v)
    if kind == "=":
        if v < 0.001:
            return "S"
        if v < 0.01:
            return "M"
        if v <= 0.05:
            return "W"
        return "NS"
    if kind in ("<", "<="):
        if v <= 0.001:
            return "S"
        if v <= 0.01:
            return "UNRES_M"
        if v <= 0.05:
            return "UNRES_W"
        return "NS"
    return "NS"  # ">" : reported as above a value


SIGNIFICANT_BANDS = {"S", "M", "W", "UNRES_M", "UNRES_W"}


def outcome_primary(label, mixed_as=0.0, flawed_as=np.nan) -> float:
    s = str(label).strip().lower() if label is not None else ""
    if s == "successful":
        return 1.0
    if s == "failed":
        return 0.0
    if s == "mixed":
        return mixed_as
    if s == "statistically successful but flawed":
        return flawed_as
    return np.nan  # 'descriptive only', 'uninformative', missing


def outcome_signal(t, v) -> float:
    """Secondary outcome: the replication reached p < .05 (direction not checked)."""
    kind = norm_type(t)
    if kind is None or v is None or (isinstance(v, float) and np.isnan(v)):
        return np.nan
    v = float(v)
    if kind == "=":
        return 1.0 if v < 0.05 else 0.0
    if kind in ("<", "<="):
        return 1.0 if v <= 0.05 else np.nan
    return 0.0 if v >= 0.05 else np.nan


def per_paper(df: pd.DataFrame, clusters: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """For each original paper: does it have rows in this group, the mean outcome of those
    rows, and (for the unweighted rate) their sum and count."""
    g = df.groupby("doi_o")["y"].agg(["sum", "count"]).reindex(clusters)
    count = g["count"].fillna(0).to_numpy()
    total = g["sum"].fillna(0).to_numpy()
    has = (count > 0).astype(float)
    mean = np.divide(total, count, out=np.zeros_like(total), where=count > 0)
    return has, mean, np.stack([total, count])


def diff_with_ci(a: pd.DataFrame, b: pd.DataFrame, rng, weighted=True) -> dict:
    """Rate(a) - rate(b), with a cluster bootstrap over original DOIs.

    Weighted (the pre-registered default): each original paper counts once in each group,
    so its replications in that group are averaged first and the group rate is the mean
    of those paper averages. Both groups are resampled from the union of papers, so a
    paper with rows in both groups moves as one unit; a paper drawn twice counts twice."""
    clusters = np.array(sorted(set(a["doi_o"]) | set(b["doi_o"])), dtype=object)
    ha, ma, sa = per_paper(a, clusters)
    hb, mb, sb = per_paper(b, clusters)

    def rates(w: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        # w: (draws, papers) multiplicities
        if weighted:
            ra = (w @ (ha * ma)) / (w @ ha)
            rb = (w @ (hb * mb)) / (w @ hb)
        else:
            ra = (w @ sa[0]) / (w @ sa[1])
            rb = (w @ sb[0]) / (w @ sb[1])
        return ra, rb

    pa, pb = rates(np.ones((1, len(clusters))))
    k = len(clusters)
    draws = rng.integers(0, k, size=(BOOT, k))
    w = np.zeros((BOOT, k))
    np.add.at(w, (np.arange(BOOT)[:, None], draws), 1.0)
    with np.errstate(invalid="ignore", divide="ignore"):
        ra, rb = rates(w)
    d = ra - rb
    d = d[np.isfinite(d)]
    lo, hi = np.percentile(d, [2.5, 97.5])
    return {"rate_a": float(pa[0]), "rate_b": float(pb[0]), "diff": float(pa[0] - pb[0]),
            "ci95": [float(lo), float(hi)], "boot_used": int(len(d)),
            "n_a": int(len(a)), "papers_a": int(a["doi_o"].nunique()),
            "n_b": int(len(b)), "papers_b": int(b["doi_o"].nunique())}


def h1_verdict(r: dict) -> str:
    d, (lo, hi) = r["diff"] * 100, (r["ci95"][0] * 100, r["ci95"][1] * 100)
    if d >= MIN_DIFF_PP and lo > 0:
        return "CUE HOLDS"
    if hi < 0:
        return "CUE REVERSED"
    return "CUE NOT SHOWN"


def h2_verdict(r: dict) -> str:
    d, lo = r["diff"] * 100, r["ci95"][0] * 100
    if lo > 0 and d >= MIN_DIFF_PP:
        return "SIZE HELPS (piece may NOT say sample size fails as a cue)"
    if lo > 0:
        return "SIZE HELPS A LITTLE (below the 10-point bar; piece must say so)"
    return "NO DETECTABLE ADVANTAGE FOR BIG SAMPLES (piece may say size is not a cue on its own)"


def pct(x: float) -> str:
    return f"{x * 100:.0f}%"


def pp(x: float) -> str:
    return f"{x * 100:+.0f}"


def load(dummy: bool) -> pd.DataFrame:
    raw = open(DATA, "rb").read()
    got = hashlib.sha256(raw).hexdigest()
    if got != DATA_SHA256:
        print(f"DATA HASH MISMATCH: {got} != pinned {DATA_SHA256}")
        sys.exit(1)
    cols = pd.read_csv(DATA, nrows=0).columns.tolist()
    use = [c for c in cols if not (dummy and c in OUTCOME_COLS)]
    df = pd.read_csv(DATA, usecols=use)
    df["band"] = [p_band(t, v) for t, v in zip(df["pval_type_o"], df["pval_value_o"])]
    return df


def code_outcome(df: pd.DataFrame, dummy: bool, rng, **kw) -> pd.DataFrame:
    out = df.copy()
    if dummy and "--plant" in sys.argv:
        # positive control: S replicates at 65%, everything else at 45%; big samples no better
        p = np.where(out["band"] == "S", 0.65, 0.45)
        out["y"] = (rng.random(len(out)) < p).astype(float)
    elif dummy:
        out["y"] = rng.integers(0, 2, size=len(out)).astype(float)
    else:
        out["y"] = [outcome_primary(x, **kw) for x in out["reported_success"]]
    return out.dropna(subset=["y"])


def code_signal(df: pd.DataFrame, dummy: bool, rng) -> pd.DataFrame:
    out = df.copy()
    if dummy:
        out["y"] = rng.integers(0, 2, size=len(out)).astype(float)
    else:
        out["y"] = [outcome_signal(t, v) for t, v in zip(out["pval_type_r"], out["pval_value_r"])]
    return out.dropna(subset=["y"])


def run_h1(df: pd.DataFrame, rng, weighted=True, w_bands=("W",)) -> dict:
    s = df[df["band"] == "S"]
    w = df[df["band"].isin(w_bands)]
    return diff_with_ci(s, w, rng, weighted=weighted)


def run_h2(df: pd.DataFrame, rng) -> dict:
    d = df[df["band"].isin(SIGNIFICANT_BANDS) & df["n_o"].notna()]
    q1, q2 = d["n_o"].quantile([1 / 3, 2 / 3]).tolist()
    top, bottom = d[d["n_o"] > q2], d[d["n_o"] <= q1]
    r = diff_with_ci(top, bottom, rng)
    r["cuts"] = [float(q1), float(q2)]
    r["spearman_n_vs_outcome"] = float(d["n_o"].rank().corr(d["y"].rank()))
    r["n_all"] = int(len(d))
    return r


def subsets(df: pd.DataFrame) -> dict:
    return {
        "all": df,
        "excluding OSC 2015": df[~df["doi_r"].astype(str).str.contains(OSC_DOI, case=False)],
        "no author overlap": df[df["author_overlap"] == 0],
    }


def main() -> int:
    dummy = "--dummy" in sys.argv
    rng = np.random.default_rng(SEED)
    base = load(dummy)
    res = {"dummy": dummy, "data_sha256": DATA_SHA256, "seed": SEED, "boot": BOOT,
           "band_counts_all_rows": base["band"].value_counts().to_dict()}
    tag = "DUMMY RUN (random outcomes; no outcome column read)" if dummy else "REAL RUN"
    print(f"FReD two-cue test · {tag}")
    print(f"data {os.path.basename(DATA)} sha256 {DATA_SHA256[:12]}… · rows {len(base)} · seed {SEED} · {BOOT} resamples\n")

    primary = code_outcome(base, dummy, rng)
    res["h1"], res["h2"] = {}, {}
    for name, d in subsets(primary).items():
        r1 = run_h1(d, rng)
        r1["verdict"] = h1_verdict(r1)
        res["h1"][name] = r1
        r2 = run_h2(d, rng)
        r2["verdict"] = h2_verdict(r2)
        res["h2"][name] = r2

    a = res["h1"]["all"]
    print(f"HEADLINE H1: original p < .001 replicated {pct(a['rate_a'])} ({a['n_a']} replications, {a['papers_a']} papers); "
          f".01 to .05 replicated {pct(a['rate_b'])} ({a['n_b']}, {a['papers_b']} papers); "
          f"difference {pp(a['diff'])} points [95% CI {pp(a['ci95'][0])} to {pp(a['ci95'][1])}] → {a['verdict']}")
    b = res["h2"]["all"]
    print(f"HEADLINE H2: largest third of original samples (n > {b['cuts'][1]:.0f}) replicated {pct(b['rate_a'])}; "
          f"smallest third (n ≤ {b['cuts'][0]:.0f}) replicated {pct(b['rate_b'])}; difference {pp(b['diff'])} points "
          f"[95% CI {pp(b['ci95'][0])} to {pp(b['ci95'][1])}]; Spearman {b['spearman_n_vs_outcome']:+.2f} → {b['verdict']}\n")

    print("Pre-registered subsets (same rules):")
    for name in ("excluding OSC 2015", "no author overlap"):
        a, b = res["h1"][name], res["h2"][name]
        print(f"  {name:20} H1 {pct(a['rate_a'])} vs {pct(a['rate_b'])} ({pp(a['diff'])} [{pp(a['ci95'][0])}, {pp(a['ci95'][1])}]) {a['verdict']}"
              f" · H2 {pct(b['rate_a'])} vs {pct(b['rate_b'])} ({pp(b['diff'])} [{pp(b['ci95'][0])}, {pp(b['ci95'][1])}])")

    print("\nPre-registered sensitivity checks on H1 (all rows):")
    sens = {}
    sens["mixed excluded"] = run_h1(code_outcome(base, dummy, rng, mixed_as=np.nan), rng)
    sens["flawed counted as success"] = run_h1(code_outcome(base, dummy, rng, flawed_as=1.0), rng)
    sens["'below .05' reports added to W"] = run_h1(primary, rng, w_bands=("W", "UNRES_W"))
    sens["unweighted (each replication counts once)"] = run_h1(primary, rng, weighted=False)
    sens["secondary outcome: replication p < .05"] = run_h1(code_signal(base, dummy, rng), rng)
    for k, r in sens.items():
        r["verdict"] = h1_verdict(r)
        print(f"  {k:44} {pct(r['rate_a'])} vs {pct(r['rate_b'])} ({pp(r['diff'])} [{pp(r['ci95'][0])}, {pp(r['ci95'][1])}]) {r['verdict']}")
    res["h1_sensitivity"] = sens

    out = RESULTS.replace(".json", ".dummy.json") if dummy else RESULTS
    json.dump(res, open(out, "w"), indent=2, default=float)
    print(f"\nwrote {os.path.relpath(out, HERE)}")
    return 0


if __name__ == "__main__":
    sys.exit(main())