#!/usr/bin/env python3
"""smallgroup_check.py — the pre-registered small-group check of the Re-check card's sign 4 (survives-a-recheck, Stage 3).

The held-out test judged sign 4 ("a claim most people would doubt") by the doubt of about 125 lay readers per finding.
A reader running the card has themselves and a few people to ask. This re-draws sign 4 from k = 1, 3 or 5 of
Hoogeveen et al.'s own lay raters (description-only condition, same exclusions as the held-out test) and asks whether
the frozen card still sorts the 27 held-out findings. Signs 1 (weak pass) and 2 (interaction) stay as frozen.

  primary      panel draws: the same k raters judge all 27 findings (raters who answered all 27), as a reader asking
               the same few people would
  sensitivity  independent draws: k raters drawn afresh for each finding from everyone who answered it

Scored by the rules in `PREREG — survives-a-recheck — Small-group check — <date>.md`.

    python3 smallgroup_check.py            # real outcomes
    python3 smallgroup_check.py --dummy    # outcomes permuted (control: the sort should vanish)
    python3 smallgroup_check.py --plant    # outcome := the crowd's majority call (control: the sort should appear)
"""
from __future__ import annotations

import argparse
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", "hoogeveen2020")
SIGNS = os.path.join(DATA, "heldout_signs.csv")
STUDIES = os.path.join(DATA, "DescriptionAndStatisticsPerStudy.csv")
SHARES = os.path.join(DATA, "lay_shares_description_only.csv")
RAW = ["Data_MTurk_AN.csv", "Data_Social_Media_AN.csv", "Data_UvA_Lab_AN.csv"]
SHA = {
    "heldout_signs.csv": "0875008358e65f7476a769b45fbaed0717023d8853206a89cdeaea268ba28eb4",
    "DescriptionAndStatisticsPerStudy.csv": "7fba68c006acd4f781d9c4dbc7045fdac986a7df662dbb9160bd0e1661240f40",
}
KS = (1, 3, 5)
DRAWS = 10_000
SEED = 20261027
RULE_GAP = 0.20          # the held-out test's own per-draw rule: act - wait >= 20 points and act >= try >= wait
PASS_SHARE = 0.75        # k keeps the sort if the per-draw rule holds in at least 75% of draws ...
MEDIAN_GAP = 0.20        # ... and the median act - wait gap is at least 20 points


def sha(p: str) -> str:
    return hashlib.sha256(open(p, "rb").read()).hexdigest()


def raters() -> tuple[np.ndarray, pd.DataFrame]:
    """Per-rater predictions, description-only: 1 = will replicate, 0 = will not, NaN = no usable answer."""
    rows = []
    for f in RAW:
        d = pd.read_csv(os.path.join(DATA, f), low_memory=False)
        d = d[(d["Finished"] == 1) & (d["phd"] == 2) & (d["knowledge"] == 5) & (d["Conditie"] == 1)]
        for _, r in d.iterrows():
            ans = []
            for s in range(1, 28):
                v = np.nan
                for suf in ("", "NL"):
                    b, u = f"{s}_des.belief{suf}", f"{s}_des.0{suf}"
                    if b in d.columns and pd.notna(r[b]) and pd.isna(r.get(u, np.nan)):
                        v = float(r[b] == 1)
                ans.append(v)
            rows.append(ans)
    m = np.array(rows)
    share = np.nanmean(m, axis=0)
    n = np.sum(~np.isnan(m), axis=0)
    frozen = pd.read_csv(SHARES).sort_values("study")
    return m, pd.DataFrame({"share": share, "n": n, "frozen_share": frozen["share_predict_replicate"].values,
                            "frozen_n": frozen["n_raters"].values})


def score(doubt: np.ndarray, weak: np.ndarray, inter: np.ndarray, held: np.ndarray) -> dict:
    """doubt: draws x 27 booleans. Returns per-draw rates and the pooled summaries."""
    signs = weak[None, :].astype(int) + inter[None, :].astype(int) + doubt.astype(int)
    act, tr, wait = signs == 0, signs == 1, signs >= 2
    out = {}
    for name, mask in (("act", act), ("try small", tr), ("wait", wait)):
        n = mask.sum(1)
        k = (mask & held[None, :]).sum(1)
        out[name] = (k, n)
    with np.errstate(invalid="ignore", divide="ignore"):
        ra, rt, rw = (out[a][0] / out[a][1] for a in ("act", "try small", "wait"))
    both = (out["act"][1] > 0) & (out["wait"][1] > 0)
    gap = np.where(both, ra - rw, np.nan)
    mono = both & (out["try small"][1] > 0) & (ra >= rt) & (rt >= rw)
    rule = mono & (gap >= RULE_GAP)
    decided = act | wait
    correct = (act & held[None, :]) | (wait & ~held[None, :])
    lay_alone = (~doubt) == held[None, :]
    return {
        "draws": int(len(doubt)),
        "both_cells_nonempty": float(both.mean()),
        "median_gap": float(np.nanmedian(gap)) if both.any() else None,
        "gap_2.5": float(np.nanpercentile(gap, 2.5)) if both.any() else None,
        "gap_97.5": float(np.nanpercentile(gap, 97.5)) if both.any() else None,
        "rule_share": float(rule.mean()),
        "pooled": {a: {"held": int(out[a][0].sum()), "n": int(out[a][1].sum()),
                       "rate": float(out[a][0].sum() / out[a][1].sum()) if out[a][1].sum() else None,
                       "mean_size": float(out[a][1].mean())} for a in out},
        "decided_accuracy": float(correct.sum() / decided.sum()) if decided.sum() else None,
        "decided_coverage": float(decided.mean()),
        "lay_alone_accuracy": float(lay_alone.mean()),
    }


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--dummy", action="store_true")
    ap.add_argument("--plant", action="store_true")
    a = ap.parse_args()
    for f, h in SHA.items():
        if sha(os.path.join(DATA, f)) != h:
            print(f"HASH MISMATCH: {f} changed")
            return 1
    m, chk = raters()
    if not (np.allclose(chk["share"], chk["frozen_share"], atol=1e-12) and (chk["n"] == chk["frozen_n"]).all()):
        print("REPRODUCTION FAILED: the rebuilt raters do not give the held-out test's frozen shares")
        return 1
    signs = pd.read_csv(SIGNS).sort_values("study")
    weak = signs["weak_pass"].values.astype(bool)
    inter = signs["interaction"].values.astype(bool)
    crowd_doubt = signs["doubted"].values.astype(bool)
    if a.dummy or a.plant:
        if a.plant:
            held = ~crowd_doubt
            label = "PLANTED (outcome = the crowd's majority call)"
        else:
            held = np.random.default_rng(7).permutation(signs["study"].values % 2 == 0)
            label = "DUMMY (outcomes permuted)"
    else:
        out = pd.read_csv(STUDIES)[["Study number", "RS outcome"]].rename(columns={"Study number": "study"})
        held = signs.merge(out, on="study", how="left")["RS outcome"].astype(int).values.astype(bool)
        if not held[signs["study"].values == 10][0]:   # coding check fixed in the held-out prereg (Karpicke & Blunt held)
            held = ~held
            print("coding check: study 10 read 0, so the outcome coding was reversed (as pre-registered)")
        label = "REAL outcomes"
    print(f"Small-group check · {label} · {m.shape[0]} raters rebuilt, shares reproduce the held-out test exactly")

    complete = m[~np.isnan(m).any(1)]
    print(f"panel pool: {len(complete)} raters who answered all 27 · draws per k: {DRAWS}\n")
    res = {"label": label, "n_raters": int(m.shape[0]), "n_complete": int(len(complete)), "by_k": {}}
    for k in KS:
        rng = np.random.default_rng(SEED + k)
        # primary: the same k raters judge every finding
        idx = np.array([rng.choice(len(complete), size=k, replace=False) for _ in range(DRAWS)])
        votes = complete[idx]                                   # draws x k x 27
        panel_doubt = votes.sum(1) < k / 2                      # fewer than half predict it replicates
        # sensitivity: k raters drawn afresh per finding from that finding's answerers
        ind_doubt = np.zeros((DRAWS, 27), dtype=bool)
        for s in range(27):
            col = m[:, s][~np.isnan(m[:, s])]
            pick = np.array([rng.choice(len(col), size=k, replace=False) for _ in range(DRAWS)])
            ind_doubt[:, s] = col[pick].sum(1) < k / 2
        r = {}
        for design, doubt in (("panel", panel_doubt), ("independent", ind_doubt)):
            sc = score(doubt, weak, inter, held)
            sc["agreement_with_crowd"] = float((doubt == crowd_doubt[None, :]).mean())
            r[design] = sc
        p = r["panel"]
        keeps = p["median_gap"] is not None and p["median_gap"] >= MEDIAN_GAP and p["rule_share"] >= PASS_SHARE
        r["verdict"] = f"SMALL GROUP KEEPS THE SORT (k = {k})" if keeps else f"NOT KEPT AT k = {k}"
        res["by_k"][str(k)] = r
        for design in ("panel", "independent"):
            s = r[design]
            pr = s["pooled"]
            print(f"k = {k} {design:>11}: agree w/ crowd {s['agreement_with_crowd'] * 100:.0f}% · "
                  f"act {pr['act']['rate'] * 100:.0f}% (avg {pr['act']['mean_size']:.1f}) · "
                  f"try {pr['try small']['rate'] * 100:.0f}% · wait {pr['wait']['rate'] * 100:.0f}% "
                  f"(avg {pr['wait']['mean_size']:.1f}) · median gap {s['median_gap'] * 100:+.0f} "
                  f"[{s['gap_2.5'] * 100:+.0f}, {s['gap_97.5'] * 100:+.0f}] · rule holds {s['rule_share'] * 100:.0f}% · "
                  f"decided {s['decided_accuracy'] * 100:.0f}% right on {s['decided_coverage'] * 100:.0f}% · "
                  f"lay alone {s['lay_alone_accuracy'] * 100:.0f}%")
        print(f"        -> {r['verdict']}\n")
    passing = [k for k in KS if res["by_k"][str(k)]["verdict"].startswith("SMALL GROUP")]
    res["smallest_passing_k"] = passing[0] if passing else None
    print(f"smallest k that keeps the sort: {res['smallest_passing_k']}")
    if not (a.dummy or a.plant):
        dst = os.path.join(HERE, "smallgroup-check-results.json")
    else:
        dst = os.path.join(HERE, f"smallgroup-check-results.{'plant' if a.plant else 'dummy'}.json")
    json.dump(res, open(dst, "w"), indent=2, default=float)
    print(f"wrote {os.path.basename(dst)}")
    return 0


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