checking which cohort users have no mentor on record

import os, re, gzip, json, pickle
from pathlib import Path
import pandas as pd
import numpy as np

# ---- paths ----
ROOT      = Path("/home/yubozhou/2026_summer/wikipedia_2sls/2sls_pipeline")
DATASET   = ROOT / "inputs/z_mentorship_state.tsv"                       # Martin's 4.98M list (userId, mentorshipState)
SNAPSHOT  = ROOT / "inputs/cov_mentor_mentee_assignment_20260530.sql.gz" # mentor-mentee snapshot dump
REG_FILE  = ROOT / "build/registrations.jsonl"                           # 11.87M registration timestamps
CLAIMLOG  = ROOT / "inputs/excl_mentor_claim_log_raw.jsonl"              # 1.59M claim/set mentor events

# ---- cache dir ----
CACHE = ROOT / "analysis/diagnose_missing_mentors/cache"
CACHE.mkdir(parents=True, exist_ok=True)

# ---- rollout phases (a user's reg date decides their phase) ----
# Dates = gerrit committer dates = production effective dates.
# Early phases are NESTED inside the Newcomer Homepage rollout (homepage < 100%),
# so they are NOT clean treatment periods. The per-user mentorship rollout only
# becomes identifying once Homepage = 100% on 2022-03-07.
# Columns: (label, start, end, mentorship_pct, homepage_pct, identifying)
ROLLOUT = [
    ("pre_anything",                  "2021-01-01", "2021-06-07", None, 0.00, False),
    ("homepage_2pct",                 "2021-06-08", "2021-09-19", None, 0.02, False),
    ("homepage_25pct_mentor_20pct",   "2021-09-20", "2022-03-06", 0.20, 0.25, False),
    ("p10",                           "2022-03-07", "2023-07-10", 0.10, 1.00, True),
    ("p25",                           "2023-07-11", "2023-10-04", 0.25, 1.00, True),
    ("p50",                           "2023-10-05", "2025-02-02", 0.50, 1.00, True),
    ("p75",                           "2025-02-03", "2025-02-16", 0.75, 1.00, True),
    ("p100",                          "2025-02-17", "2026-6-8", 1.00, 1.00, False),  # no control arm
]

def helper_cache(name):
    return CACHE / name

# load dataset
ds = pd.read_csv(DATASET, sep="\t", dtype={"userId": "int64", "mentorshipState": "string"})
print(len(ds), "users in dataset")
print(ds["mentorshipState"].value_counts(dropna=False), "\n")

# parse snapshot (cached): mentee_id -> set of roles
snap_pkl = helper_cache("snap.pkl")
if snap_pkl.exists():
    snap = pickle.load(open(snap_pkl, "rb"))
else:
    pat = re.compile(rb"\((\d+),'([^']+)',(\d+),(\d+)\)")
    snap = {}
    with gzip.open(SNAPSHOT, "rb") as f:
        for line in f:
            if not line.startswith(b"INSERT"):
                continue
            for m in pat.finditer(line):
                snap.setdefault(int(m.group(1)), set()).add(m.group(2).decode())
    pickle.dump(snap, open(snap_pkl, "wb"))
print(len(snap), "mentees have at least one snapshot row\n")

# mismatch: per mentorshipState, how many users have NO row in snapshot
ds["has_snapshot_row"] = ds["userId"].map(lambda u: u in snap)
g = ds.groupby("mentorshipState")["has_snapshot_row"].agg(n="size", in_snap="sum")
g["no_row"]     = g["n"] - g["in_snap"]
g["no_row_pct"] = (g["no_row"] / g["n"] * 100).round(4)
print(g[["n", "in_snap", "no_row", "no_row_pct"]])

# snapshot (cached)
snap_pkl = helper_cache("snap.pkl")
if snap_pkl.exists():
    snap = pickle.load(open(snap_pkl, "rb"))
else:
    pat = re.compile(rb"\((\d+),'([^']+)',(\d+),(\d+)\)")
    snap = {}
    with gzip.open(SNAPSHOT, "rb") as f:
        for line in f:
            if not line.startswith(b"INSERT"):
                continue
            for m in pat.finditer(line):
                snap.setdefault(int(m.group(1)), set()).add(m.group(2).decode())
    pickle.dump(snap, open(snap_pkl, "wb"))

# registration timestamps (cached)
reg_pkl = helper_cache("reg_ts.pkl")
if reg_pkl.exists():
    reg = pickle.load(open(reg_pkl, "rb"))
else:
    reg = {}
    with open(REG_FILE) as f:
        for line in f:
            o = json.loads(line)
            reg[o["uid"]] = o["reg_ts"]
    pickle.dump(reg, open(reg_pkl, "wb"))

# attach month + no_row flag
ds["reg_ts"]    = pd.to_datetime(ds["userId"].map(reg), format="mixed", errors="coerce")
ds["reg_month"] = ds["reg_ts"].dt.to_period("M").astype("string")
ds["no_row"]    = ~ds["userId"].map(lambda u: u in snap)

# table 1: counts of each state per month
counts = ds.pivot_table(index="reg_month", columns="mentorshipState",
                        aggfunc="size", fill_value=0)
counts["total"] = counts.sum(axis=1)
print("=== counts per month ===")
print(counts.to_string())

# table 2: per month, no_row count of each state
norow = (ds[ds["no_row"]]
        .pivot_table(index="reg_month", columns="mentorshipState",
                    aggfunc="size", fill_value=0)
        .reindex(counts.index, fill_value=0))
print("\n=== no_row count per month ===")
print(norow.to_string())

# table 3: per month, no_row count of each state / that month's total registered users (%)
norow = (ds[ds["no_row"]]
        .pivot_table(index="reg_month", columns="mentorshipState",
                    aggfunc="size", fill_value=0)
        .reindex(counts.index, fill_value=0))
pct = norow.div(counts["total"], axis=0).mul(100).round(4)
print("\n=== no_row count as pct of monthly total ===")
print(pct.to_string())

# =============================================================================
# WHY A USER CAN BE "no_row" (in Martin's list but has NO mentor row in snapshot)
#
# Scope: this runs on ALL no_row users across every mentorshipState
#        (0 / 50 / unset), not just state=0. `nr = ds[ds["no_row"]]` has no
#        state filter, so the totals below are the 0 + 50 + unset combined.
#
# Assignment mechanism differs by state:
#   - unset / 0 : a mentor is assigned synchronously at registration
#                 (MentorHooks::onLocalUserCreated). state=0 (DISABLED) only
#                 hides the UI; it does NOT block the backend assignment.
#   - 50        : proactive-assignment flag (2024-era). Assignment is done by an
#                 async backend job, NOT onLocalUserCreated. So a no_row state=50
#                 user can simply mean that async job never ran / failed — a
#                 different mechanism from the reasons below.
#
# A user ends up with no mentor row for one of FIVE reasons. These are NOT
# mutually exclusive (e.g. a user can be both pre-rollout and autocreated); the
# final per-user classification in the later cell resolves overlaps by priority.
#
#   1. pre-rollout  : registered before en.wiki turned mentorship on (~2021-06).
#                     The feature wasn't live yet, so nobody was assigned.
#                     -> detectable from reg_ts.                        [we HAVE this]
#   2. autocreated  : the en.wiki local account was auto-created (CentralAuth) or
#                     created by someone else, not self-registered. onLocalUserCreated
#                     returns early for these.
#                     -> detectable from is_self == False.              [we HAVE this]
#   3. blocked      : user was indefinitely blocked, which DROPS the mentor row.
#                     -> from block_status.jsonl (fetched in a later cell).
#                                                                       [we NOW HAVE this]
#   4. no mentor available at registration : the auto-assign mentor pool was empty
#                     / everyone excluded at that moment. (There may be other
#                     causes I cannot enumerate.)
#                     -> no direct evidence; only by elimination.       [residual]
#   5. deleted      : the account was deleted, so its registration record is gone
#                     and reg_ts is missing.
#                     -> detectable from reg_ts being NaN.              [we HAVE this]
#
# This cell quantifies reasons 1, 2 and 5 among the no_row users (the ones we can
# cleanly identify here), and reports how many remain for reasons 3 and 4 (which
# the block-fetch cell resolves later).
# =============================================================================

# ---- load dataset ----
if "ds" not in globals():
    ds = pd.read_csv(DATASET, sep="\t", dtype={"userId": "int64", "mentorshipState": "string"})

# ---- snapshot (cached): used to flag no_row ----
snap_pkl = helper_cache("snap.pkl")
if snap_pkl.exists():
    snap = pickle.load(open(snap_pkl, "rb"))
else:
    pat = re.compile(rb"\((\d+),'([^']+)',(\d+),(\d+)\)")
    snap = {}
    with gzip.open(SNAPSHOT, "rb") as f:
        for line in f:
            if not line.startswith(b"INSERT"):
                continue
            for m in pat.finditer(line):
                snap.setdefault(int(m.group(1)), set()).add(m.group(2).decode())
    pickle.dump(snap, open(snap_pkl, "wb"))

# ---- registration info (cached): reg_ts + is_self in one pass ----
reginfo_pkl = helper_cache("reg_info.pkl")
if reginfo_pkl.exists():
    reg_ts_map, is_self_map = pickle.load(open(reginfo_pkl, "rb"))
else:
    reg_ts_map, is_self_map = {}, {}
    with open(REG_FILE) as f:
        for line in f:
            o = json.loads(line)
            reg_ts_map[o["uid"]]  = o["reg_ts"]
            is_self_map[o["uid"]] = bool(o.get("is_self"))
    pickle.dump((reg_ts_map, is_self_map), open(reginfo_pkl, "wb"))

# ---- tag every user with no_row / reg_ts / is_self ----
ds["no_row"]  = ~ds["userId"].map(lambda u: u in snap)
ds["reg_ts"]  = pd.to_datetime(ds["userId"].map(reg_ts_map), format="mixed", errors="coerce")
ds["is_self"] = ds["userId"].map(is_self_map)   # True=self-registered, False=autocreated/created-by-other, NaN=unknown (deleted account)

# ---- restrict to no_row users ----
ROLLOUT_LIVE = pd.Timestamp("2021-06-01")   # en.wiki mentorship becomes effective ~here
nr = ds[ds["no_row"]].copy()

nr["cause_pre_rollout"] = nr["reg_ts"] < ROLLOUT_LIVE   # reason 1
nr["cause_autocreated"] = nr["is_self"] == False        # reason 2
nr["reg_ts_missing"]    = nr["reg_ts"].isna()           # deleted accounts (no reg_ts)

total = len(nr)
print(f"total no_row users: {total:,}\n")

n1 = nr['cause_pre_rollout'].sum()
n2 = nr['cause_autocreated'].sum()
n3 = nr['reg_ts_missing'].sum()
print(f"reason 1  pre-rollout (reg_ts < {ROLLOUT_LIVE.date()}): {n1:,}  ({n1/total*100:.2f}%)")
print(f"reason 2  autocreated (is_self == False)             : {n2:,}  ({n2/total*100:.2f}%)")
print(f"(aside)   reg_ts unknown (deleted account)           : {n3:,}  ({n3/total*100:.2f}%)")

print("\noverlap (a user may match both):")
print(pd.crosstab(nr["cause_pre_rollout"], nr["cause_autocreated"],
                rownames=["pre_rollout"], colnames=["autocreated"]))

rest = nr[~nr["cause_pre_rollout"] & ~nr["cause_autocreated"]]
print(f"\nremaining no_row after excluding reasons 1 & 2: {len(rest):,}({len(rest)/total*100:.2f}%)")
print("this remainder needs block data (reason 3) and elimination (reason 4).")

# export the remainder (post-rollout, self-registered, no_row) with usernames
need = set(ds.loc[ds["no_row"] & (ds["reg_ts"] >= ROLLOUT_LIVE) & (ds["is_self"] != False),
                "userId"].astype(int))
print("remainder:", len(need))

uid2name = {}
with open(REG_FILE) as f:
    for line in f:
        o = json.loads(line)
        if o["uid"] in need:
            uid2name[o["uid"]] = o["name"]
print("names resolved:", len(uid2name))

out = ROOT / "analysis/diagnose_missing_mentors/remainder_to_block.tsv"
with open(out, "w") as fo:
    for u, n in uid2name.items():
        fo.write(f"{u}\t{n}\n")
print("written:", out)

import time
from urllib.parse import urlencode
from urllib.request import Request, urlopen
from urllib.error import HTTPError, URLError

API = "https://en.wikipedia.org/w/api.php"
UA  = "WikiMentorResearch/1.0 (academic; contact: yubozhou@umich.edu)"
INP = ROOT / "analysis/diagnose_missing_mentors/remainder_to_block.tsv"
OUT = ROOT / "analysis/diagnose_missing_mentors/block_status.jsonl"

def api_get(params):
    url = API + "?" + urlencode(dict(params, format="json"))
    for attempt in range(6):
        try:
            req = Request(url, headers={"User-Agent": UA})
            with urlopen(req, timeout=60) as r:
                return json.loads(r.read().decode())
        except (HTTPError, URLError, TimeoutError) as e:
            time.sleep(0.4 * (2 ** attempt))
    raise RuntimeError("API failed: " + url)

# load the list of (uid, name)
todo = []
with open(INP) as f:
    for line in f:
        uid, name = line.rstrip("\n").split("\t", 1)
        todo.append((int(uid), name))

# resume: skip uids already fetched
done = set()
if OUT.exists():
    with open(OUT) as f:
        for line in f:
            try: done.add(json.loads(line)["uid"])
            except Exception: pass
todo = [(u, n) for u, n in todo if u not in done]
print(f"to fetch: {len(todo):,} (already done {len(done):,})")

# fetch in batches of 50
name2uid = {n: u for u, n in todo}
with open(OUT, "a") as fout:
    for i in range(0, len(todo), 50):
        names = [n for _, n in todo[i:i+50]]
        data = api_get({"action": "query", "list": "users",
                        "ususers": "|".join(names), "usprop": "blockinfo"})
        seen = set()
        for u in data.get("query", {}).get("users", []):
            name = u.get("name"); uid = name2uid.get(name)
            if uid is None: continue
            blocked = "blockid" in u
            expiry  = u.get("blockexpiry")
            indef   = blocked and expiry in ("infinity","infinite","indefinite","never")
            fout.write(json.dumps({"uid": uid, "name": name, "blocked": blocked,
                                    "indefinite": bool(indef), "expiry": expiry,
                                    "reason": u.get("blockreason")}) + "\n")
            seen.add(name)
        for name in names:
            if name not in seen:
                fout.write(json.dumps({"uid": name2uid[name], "name": name, "blocked": False,
                                        "indefinite": False, "expiry": None,
                                        "reason": "api_no_return"}) + "\n")
        fout.flush()
        if (i // 50) % 20 == 0: print(f"  {i+len(names):,}/{len(todo):,}")
        time.sleep(0.4)
print("fetch done.")

# ---- results: reason 3 vs reason 4 ----
blk = pd.read_json(OUT, lines=True)
indef_uids = set(blk.loc[blk["indefinite"], "uid"])

rest = ds[ds["no_row"] & (ds["reg_ts"] >= ROLLOUT_LIVE) & (ds["is_self"] != False)].copy()
rest["cause_blocked"] = rest["userId"].isin(indef_uids)
n_rest = len(rest)
n_blk  = int(rest["cause_blocked"].sum())
print(f"remainder (reasons 3 & 4): {n_rest:,}")
print(f"reason 3  indefinitely blocked (row dropped): {n_blk:,}  ({n_blk/n_rest*100:.2f}%)")
print(f"reason 4  no mentor avail / unknown residual: {n_rest-n_blk:,}({(n_rest-n_blk)/n_rest*100:.2f}%)")

# See if any of those who didn't have a mentor row (no_row) had actually ASKED A QUESTION (and thus were more likely to be impacted by missing mentorship, if they were in the treatment group).

# ---- 1. all mentee usernames who asked a question ----
QFILE = ROOT / "inputs/d_mentor_questions.jsonl"
askers = set()
with open(QFILE) as f:
    for line in f:
        o = json.loads(line)
        m = o.get("mentee")
        if m:
            askers.add(m)
print(f"unique mentees who asked a question: {len(askers):,}")

# ---- 2. map those usernames -> userId (stream registrations.jsonl) ----
name2uid = {}
with open(REG_FILE) as f:
    for line in f:
        o = json.loads(line)
        if o["name"] in askers:
            name2uid[o["name"]] = o["uid"]
print(f"resolved to userId: {len(name2uid):,} / {len(askers):,}")

# ---- 3. look up their state + no_row status in the dataset ----
uid2state = dict(zip(ds["userId"].astype(int), ds["mentorshipState"]))

rows = []
for name, uid in name2uid.items():
    if uid not in uid2state:
        continue                      # not in Martin's list
    state  = uid2state[uid]
    no_row = uid not in snap          # no mentor row in the current snapshot
    rows.append((uid, name, state, no_row))

q = pd.DataFrame(rows, columns=["uid", "name", "mentorshipState", "no_row"])
print(f"\naskers present in dataset: {len(q):,}")

# ---- 4. key result: asked a question BUT currently no_row, split by state ----
hit = q[q["no_row"]]
print(f"\n>>> asked a question BUT currently no mentor row (no_row): {len(hit):,}")
print(hit.groupby("mentorshipState").size().to_string())

# save for inspection
hit.to_csv(ROOT / "analysis/diagnose_missing_mentors/asked_but_no_row.tsv",
            sep="\t", index=False)
print("\nsaved -> analysis/diagnose_missing_mentors/asked_but_no_row.tsv")

# See why are those people missing mentor rows, and how many of them are indefinitely blocked (reason 3) vs potentially having no mentor available (reason 4, by elimination).


# indefinitely-blocked uids from the block fetch
OUT = ROOT / "analysis/diagnose_missing_mentors/block_status.jsonl"
blk = pd.read_json(OUT, lines=True)
indef_uids = set(blk.loc[blk["indefinite"], "uid"])

# bring reg_ts / is_self onto the 385 askers-without-row
h = hit.merge(ds[["userId", "reg_ts", "is_self"]],
            left_on="uid", right_on="userId", how="left")

def classify(r):
    if pd.isna(r["reg_ts"]):                       return "5_deleted"
    if r["is_self"] == False:                      return "2_autocreated"
    if r["reg_ts"] < pd.Timestamp(ROLLOUT_LIVE):   return "1_pre_rollout"
    if r["uid"] in indef_uids:                      return "3_indef_blocked"
    return "4_residual"

h["reason"] = h.apply(classify, axis=1)

print(f"asked a question but currently no_row: {len(h):,}\n")
print("by reason:")
print(h.groupby("reason").size().to_string())
print("\nby reason x state:")
print(pd.crosstab(h["reason"], h["mentorshipState"]))

h.to_csv(ROOT / "analysis/diagnose_missing_mentors/asked_but_no_row_classified.tsv",
        sep="\t", index=False)
print("\nsaved -> asked_but_no_row_classified.tsv")

import matplotlib.pyplot as plt

# block reasons for the indefinitely-blocked subset
blk = pd.read_json(ROOT / "analysis/diagnose_missing_mentors/block_status.jsonl", lines=True)
ib = h[h["reason"] == "3_indef_blocked"].merge(
        blk[["uid", "reason"]].rename(columns={"reason": "block_reason"}),
        on="uid", how="left")
print(f"indefinitely blocked: {len(ib)}")

# ---- bucket the messy free-text / template block reasons ----
def bucket(text):
    t = (text or "").lower()
    if "checkuser" in t:                                            return "checkuser"
    if "sock" in t:                                                 return "sockpuppet"
    if "spam" in t or "advertis" in t or "promot" in t or "paid" in t: return "spam/promo"
    if "username" in t or "ublock" in t:                            return "username"
    if "vandal" in t or "voa" in t:                                 return "vandalism"
    if "nothere" in t or "not here" in t:                           return "not-here"
    if "disrupt" in t:                                              return "disruption"
    if "harass" in t or "personal attack" in t or "npa" in t:       return "harassment"
    if "lta" in t or "long-term abuse" in t:                        return "LTA"
    if "block evasion" in t or "evasion" in t:                      return "block-evasion"
    if t.strip() == "":                                             return "(empty)"
    return "other"

ib["cat"] = ib["block_reason"].map(bucket)

# ---- count: category x state ----
tab = ib.pivot_table(index="cat", columns="mentorshipState",
                    aggfunc="size", fill_value=0)
tab = tab.loc[tab.sum(axis=1).sort_values(ascending=False).index]   # sort by total
print("\ncategory x state:")
print(tab.to_string())

# ---- stacked bar plot ----
ax = tab.plot(kind="bar", stacked=True, figsize=(11, 6))
ax.set_xlabel("block reason category")
ax.set_ylabel("number of users")
ax.set_title("Indefinitely-blocked no_row users (asked a question) — block reason by mentorshipState")
ax.legend(title="mentorshipState")
plt.xticks(rotation=35, ha="right")
plt.tight_layout()
plt.savefig(ROOT / "analysis/diagnose_missing_mentors/blocked_reason_by_state.png", dpi=130)
plt.show()
print("\nsaved -> blocked_reason_by_state.png")

# =============================================================================
# Classify ALL no_row users (state 0 + 50 + unset) into the 5 reasons, by state.
# Reuses what earlier cells already built:
#   - ds["no_row"], ds["reg_ts"], ds["is_self"]   (from cell-6)
#   - block_status.jsonl -> indef_uids            (from cell-9; covers all states,
#                                                  since the remainder was NOT
#                                                  filtered by state)
# Priority (high -> low):  5_deleted > 2_autocreated > 1_pre_rollout
#                          > 3_indef_blocked > 4_residual
# =============================================================================
import matplotlib.pyplot as plt

ROLLOUT_LIVE = pd.Timestamp("2021-06-01")

# ---- indefinitely-blocked uids (covers the 50 & unset remainders too) ----
blk = pd.read_json(ROOT / "analysis/diagnose_missing_mentors/block_status.jsonl",
                    lines=True)
indef_uids = set(blk.loc[blk["indefinite"], "uid"])
print(f"indefinitely-blocked uids in block_status: {len(indef_uids):,}")

# ---- all no_row users (should be 11,600 + 4 + 540,773 = 552,377) ----
nr = ds[ds["no_row"]].copy()
print(f"total no_row users: {len(nr):,}")
print(nr.groupby("mentorshipState").size().to_string(), "\n")

# ---- vectorized classification (assign low priority first, overwrite upward) ----
reason = pd.Series("4_residual", index=nr.index, dtype="object")
reason[nr["userId"].astype(int).isin(indef_uids)] = "3_indef_blocked"
reason[nr["reg_ts"] < ROLLOUT_LIVE]               = "1_pre_rollout"
reason[nr["is_self"] == False]                    = "2_autocreated"
reason[nr["reg_ts"].isna()]                       = "5_deleted"
nr["reason"] = reason

REASON_ORDER = ["1_pre_rollout", "2_autocreated", "3_indef_blocked",
                "4_residual", "5_deleted"]

# ---- reason x state table ----
tab = (nr.pivot_table(index="reason", columns="mentorshipState",
                    aggfunc="size", fill_value=0)
        .reindex(REASON_ORDER, fill_value=0))
tab["total"] = tab.sum(axis=1)
print("=== reason x mentorshipState ===")
print(tab.to_string())
print(f"\nsum check: {tab['total'].sum():,} (should equal total no_row above)")

tab.to_csv(ROOT / "analysis/diagnose_missing_mentors/no_row_reason_by_state.tsv",
            sep="\t")

# ---- plot (state 0 included; mostly so we see 50 & unset clearly) ----
plot_tab = tab.drop(columns="total")
ax = plot_tab.plot(kind="bar", stacked=True, figsize=(10, 6), logy=True)
ax.set_xlabel("reason")
ax.set_ylabel("number of no_row users (log scale)")
ax.set_title("Why no_row — reason x mentorshipState (all states)")
ax.legend(title="mentorshipState")
plt.xticks(rotation=25, ha="right")
plt.tight_layout()
plt.savefig(ROOT / "analysis/diagnose_missing_mentors/no_row_reason_by_state.png",
            dpi=130)
plt.show()
print("saved -> no_row_reason_by_state.png  +  no_row_reason_by_state.tsv")

# =============================================================================
# Re-checking the "5_deleted" bucket — are those 6,451 users actually deleted?
#
# Context: in the reason-by-state cell above we labelled 6,451 no_row users as
# "5_deleted" because their reg_ts is missing from build/registrations.jsonl
# (which is built from en.wiki's LOCAL user table). "Missing local reg" is not
# the same as "account deleted" — the account may simply be a global
# (CentralAuth) account whose home wiki is not en.wiki, in which case it has
# no local enwiki user row but the global account is alive and queryable on
# Meta.
#
# This cell:
#   1. Queries Meta's CentralAuth API (one uid at a time — guiid is not
#      batchable) and writes deleted_global_info.jsonl. Resume-safe.
#      API reference: https://www.mediawiki.org/wiki/API:Globaluserinfo
#   2. Summarises: where they registered, when, and whether any fall inside
#      the 2SLS identifying window (derived from the ROLLOUT table).
#
# NOTE: this cell reads the deleted_uids set from a cache pkl produced in an
# earlier ad-hoc step. If the pkl is missing, recompute it inline:
#     deleted = sorted(set(ds.loc[ds["no_row"] & ds["reg_ts"].isna(),
#                                 "userId"].astype(int)))
# =============================================================================
import pickle, time
from urllib.parse import urlencode
from urllib.request import Request, urlopen
from urllib.error import HTTPError, URLError
from collections import Counter

DEL  = ROOT / "analysis/diagnose_missing_mentors/cache/deleted_uids.pkl"
OUT  = ROOT / "analysis/diagnose_missing_mentors/deleted_global_info.jsonl"
META = "https://meta.wikimedia.org/w/api.php"
UA   = "WikiMentorResearch/1.0 (academic; contact: yubozhou@umich.edu)"
SLEEP = 0.4   # match cell-7's pacing

# cohort window derived from the ROLLOUT table (identifying phases only)
_ident = [(s, e) for label, s, e, mp, hp, ident in ROLLOUT if ident]
COHORT_START = min(s for s, _ in _ident)
COHORT_END   = max(e for _, e in _ident)
ROLLOUT_LIVE = "2021-06-01"   # en.wiki mentorship effective date (same as cell-4)
print(f"cohort window (from ROLLOUT identifying phases): {COHORT_START} .. {COHORT_END}")

# ---- 1) fetch (resume-safe; skips uids already in OUT) -----------------------
def fetch_one(uid):
    p = {"action":"query", "meta":"globaluserinfo", "guiid":uid, "format":"json"}
    url = META + "?" + urlencode(p)
    for k in range(5):
        try:
            req = Request(url, headers={"User-Agent": UA})
            with urlopen(req, timeout=30) as r:
                return json.loads(r.read().decode())
        except (HTTPError, URLError, TimeoutError):
            time.sleep(0.5 * (2**k))
    return {"error": "fail"}

deleted = sorted(pickle.load(open(DEL, "rb")))
done = set()
if OUT.exists():
    with open(OUT) as f:
        for line in f:
            try: done.add(json.loads(line)["uid"])
            except: pass
todo = [u for u in deleted if u not in done]
print(f"total={len(deleted):,}  already_fetched={len(done):,}  to_fetch={len(todo):,}")

if todo:
    with open(OUT, "a") as fout:
        for i, uid in enumerate(todo):
            d  = fetch_one(uid)
            gu = d.get("query", {}).get("globaluserinfo", d)
            rec = {"uid": uid, "name": gu.get("name"), "home": gu.get("home"),
                    "registration": gu.get("registration"),
                    "missing": "missing" in gu, "locked": "locked" in gu,
                    "hidden":  "hidden"  in gu,
                    "err": d.get("error") or gu.get("error")}
            fout.write(json.dumps(rec) + "\n")
            if i % 200 == 0:
                fout.flush()
                print(f"  {i:,}/{len(todo):,}")
            time.sleep(SLEEP)
    print("fetch done.")

# ---- 2) summarise ------------------------------------------------------------
rows = [json.loads(l) for l in open(OUT)]
print(f"\ntotal records: {len(rows):,}")

n_missing = sum(1 for r in rows if r.get("missing"))
n_locked  = sum(1 for r in rows if r.get("locked"))
n_hidden  = sum(1 for r in rows if r.get("hidden"))
n_err     = sum(1 for r in rows if r.get("err"))
n_no_reg  = sum(1 for r in rows if not r.get("registration"))
print(f"  globaluserinfo MISSING (account never existed globally): {n_missing:,}")
print(f"  locked (account exists but globally locked)            : {n_locked:,}")
print(f"  hidden                                                 : {n_hidden:,}")
print(f"  API errors / no registration timestamp                 : {n_err:,} / {n_no_reg:,}")

# robust date parsing (MediaWiki returns ISO 8601; coerce defensively)
reg_dt = pd.to_datetime(
    pd.Series([r.get("registration") for r in rows]),
    errors="coerce", utc=True,
).dt.tz_localize(None)
reg_date = reg_dt.dt.date.astype("string")

cs, ce, rl = pd.Timestamp(COHORT_START).date(), pd.Timestamp(COHORT_END).date(),pd.Timestamp(ROLLOUT_LIVE).date()
valid = reg_dt.notna()
n_pre   = int(((reg_dt.dt.date <  rl) & valid).sum())
n_btw   = int(((reg_dt.dt.date >= rl) & (reg_dt.dt.date <  cs) & valid).sum())
n_in    = int(((reg_dt.dt.date >= cs) & (reg_dt.dt.date <= ce) & valid).sum())
n_post  = int(((reg_dt.dt.date >  ce) & valid).sum())

print("\n--- registration vs cohort window ---")
print(f"  pre-rollout       (reg < {ROLLOUT_LIVE}) : {n_pre:,}")
print(f"  rollout..cohort   ({ROLLOUT_LIVE}..{COHORT_START}) : {n_btw:,}")
print(f"  IN cohort window  ({COHORT_START}..{COHORT_END}) : {n_in:,}")
print(f"  after cohort end  (> {COHORT_END})       : {n_post:,}")

yrs = Counter(d.year for d in reg_dt.dropna().dt.date)
print("\n--- registration year distribution ---")
for y in sorted(yrs):
    print(f"  {y}: {yrs[y]:,}")

homes = Counter(r.get("home") for r in rows if r.get("home"))
print("\n--- top 10 home wikis ---")
for w, c in homes.most_common(10):
    print(f"  {w}: {c:,}  ({c/len(rows)*100:.1f}%)")

print("\n=== conclusion ===")
print(f"None of the {len(rows):,} 'deleted' uids are actually deleted; all are queryable")
print(f"on CentralAuth. All {n_pre:,} with a known registration date registered BEFORE the")
print(f"mentorship rollout ({ROLLOUT_LIVE}), and 0 fall inside the identifying window")
print(f"({COHORT_START}..{COHORT_END}). They were tagged '5_deleted' only because the local")
print(f"en.wiki user table does not have their registration row — many are global accounts")
print(f"whose home wiki is not en.wiki ({homes.get('enwiki',0)/len(rows)*100:.1f}% are enwiki-home,")
print(f"the rest are eswiki/frwiki/etc). The '5_deleted' label is misleading — these are")
print(f"pre-rollout global accounts without a local enwiki user row.")

# =============================================================================
# Block reasons for ALL 20,974 indef-blocked no_row users, split by
# mentorshipState (0 / 50 / unset).
#
# cell-11 only bucketed the 373 indef-blocked who had also asked a question —
# too small to generalise from. This cell runs the same bucket function over
# the full block_status.jsonl set, so we can see whether state=0 vs unset have
# different abuse-reason mixes.
# =============================================================================
import matplotlib.pyplot as plt

blk = pd.read_json(ROOT / "analysis/diagnose_missing_mentors/block_status.jsonl",
                    lines=True)
blk = blk[blk["indefinite"]].copy()
print(f"indefinitely-blocked uids: {len(blk):,}")

# attach mentorshipState
uid2state = dict(zip(ds["userId"].astype(int), ds["mentorshipState"]))
blk["mentorshipState"] = blk["uid"].map(uid2state)
print("by state:")
print(blk.groupby("mentorshipState").size().to_string())

# same bucket function as cell-11
def bucket(text):
    t = (text or "").lower()
    if "checkuser" in t:                                                  return "checkuser"
    if "sock" in t:                                                       return "sockpuppet"
    if "spam" in t or "advertis" in t or "promot" in t or "paid" in t:    return "spam/promo"
    if "username" in t or "ublock" in t:                                  return "username"
    if "vandal" in t or "voa" in t:                                       return "vandalism"
    if "nothere" in t or "not here" in t:                                 return "not-here"
    if "disrupt" in t:                                                    return "disruption"
    if "harass" in t or "personal attack" in t or "npa" in t:             return "harassment"
    if "lta" in t or "long-term abuse" in t:                              return "LTA"
    if "block evasion" in t or "evasion" in t:                            return "block-evasion"
    if t.strip() == "":                                                   return "(empty)"
    return "other"

blk["cat"] = blk["reason"].map(bucket)

# absolute counts: category x state
tab = blk.pivot_table(index="cat", columns="mentorshipState",
                    aggfunc="size", fill_value=0)
tab["total"] = tab.sum(axis=1)
tab = tab.sort_values("total", ascending=False)
print("\n=== block reason category x mentorshipState (absolute counts) ===")
print(tab.to_string())

# within-state shares (column-wise %): lets you compare state-0 vs unset mixes
share = (tab.drop(columns="total")
            .div(tab.drop(columns="total").sum(axis=0), axis=1)
            .mul(100).round(2))
print("\n=== same table, as within-state share (%) ===")
print(share.to_string())

# save
tab.to_csv(ROOT / "analysis/diagnose_missing_mentors/block_reason_by_state_all.tsv", sep="\t")
share.to_csv(ROOT / "analysis/diagnose_missing_mentors/block_reason_share_by_state.tsv",sep="\t")

# stacked bar (absolute, log scale so state=50's tiny count doesn't disappear)
plot_tab = tab.drop(columns="total")
ax = plot_tab.plot(kind="bar", stacked=True, figsize=(11, 6), logy=True)
ax.set_xlabel("block reason category")
ax.set_ylabel("number of indef-blocked no_row users (log scale)")
ax.set_title("Block reasons for ALL indef-blocked no_row users, by mentorshipState")
ax.legend(title="mentorshipState")
plt.xticks(rotation=35, ha="right")
plt.tight_layout()
plt.savefig(ROOT / "analysis/diagnose_missing_mentors/block_reason_by_state_all.png",dpi=130)
plt.show()
print("\nsaved -> block_reason_by_state_all.{png,tsv}  +  block_reason_share_by_state.tsv")

# Pre-check: decompose `reg_ts.isna()` BEFORE the realized-rollout cell prints 7,378.
# 7,378 = (no_row & reg_ts NaN) + (has_row & reg_ts NaN)
in_snap_mask = ds["userId"].map(lambda u: u in snap)
n_total = int(ds["reg_ts"].isna().sum())
n_A     = int((~in_snap_mask & ds["reg_ts"].isna()).sum())
n_B     = int(( in_snap_mask & ds["reg_ts"].isna()).sum())
print(f"reg_ts NaN total           : {n_total:,}")
print(f"  group A: no_row  & NaN   : {n_A:,}   # the 5-reason 'deleted' bucket")
print(f"  group B: has_row & NaN   : {n_B:,}   # has mentor row but no local enwiki reg_ts")

# =============================================================================
# Realized rollout % per phase, with bounds for missing-reg users.
#
# What "Z=1" means in OUR mapping (per Martin's release doc):
#   Z=1 (UI enabled by rollout) = mentorshipState ∈ {unset, 50}
#   Z=0 (A/B disabled)          = mentorshipState == 0
#
# Why "nominal" needs two columns for the nested early phase:
# -----------------------------------------------------------
# In the late phases (p10/p25/p50/p75/p100), Newcomer Homepage rollout is
# already 100%, so EVERY new user is evaluated by the mentorship A/B code.
# In that case, nominal = mentorship_pct directly.
#
# But in the early phase `homepage_25pct_mentor_20pct`, only 25% of new users
# saw the Newcomer Homepage, and only those 25% went through the mentorship
# A/B evaluation. The other 75% never had the A/B code touch their property,
# so their mentorshipState stays `unset` -> they get tagged Z=1 in OUR mapping.
#
# So among ALL new registrants in this phase, the EXPECTED Z=1 share is:
#
#     expected_Z1_share
#         = (1 - homepage_pct) * 1.0                # never saw homepage -> unset -> Z=1
#         + homepage_pct * mentorship_pct           # saw homepage, A/B picked them -> Z=1
#
# For the early phase that's 0.75*1 + 0.25*0.20 = 0.80 = 80%. That 80% does
# NOT mean "80% were really shown the mentor UI" — it means "80% would carry
# Z=1 under our `unset∪50` mapping", because 3/4 of them never had the A/B
# code run at all. This phase is therefore identifying=False; we only use it
# as a sanity check.
#
# For the late phases (homepage_pct=1.0), the formula collapses to
# `mentorship_pct`, which is the actual rollout target.
# =============================================================================
ds["Z"] = ds["mentorshipState"].isin(["unset", "50"]).astype(int)

known = ds[ds["reg_ts"].notna()].copy()
miss  = ds[ds["reg_ts"].isna()].copy()
M_total = len(miss); M_z1 = int(miss["Z"].sum()); M_z0 = M_total - M_z1
print(f"users with reg_ts known   : {len(known):,}")
print(f"users with reg_ts MISSING : {M_total:,}  (Z=1: {M_z1:,}, Z=0: {M_z0:,})")

def phase_of(ts):
    if pd.isna(ts): return None
    d = ts.date().isoformat()
    for label, start, end, *_ in ROLLOUT:
        if start <= d <= end: return label
    return None
known["phase"] = known["reg_ts"].apply(phase_of)

def pct(z1, n): return round(100*z1/n, 3) if n else float("nan")

rows = []
for label, start, end, mentor_pct, hp_pct, ident in ROLLOUT:
    if mentor_pct is None: continue
    expected_z1 = (1 - hp_pct) * 1.0 + hp_pct * mentor_pct      # effective nominal under our Z mapping
    sub = known[known["phase"] == label]
    n  = len(sub); z1 = int(sub["Z"].sum())
    rows.append({
        "phase":                label,
        "start":                start,
        "end":                  end,
        "homepage_pct":         round(hp_pct*100, 1),
        "mentor_pct_of_hp":     round(mentor_pct*100, 1),   # 'rollout target' AMONG homepage viewers
        "nominal_Z1_all_users": round(expected_z1*100, 2),  # what we should see in OUR mapping
        "n_known":              n,
        "realized_baseline_%":  pct(z1, n),
        "if_all_miss_here_actual_Z": pct(z1 + M_z1,    n + M_total),
        "lower_all_miss_as_Z0":      pct(z1,           n + M_total),
        "upper_all_miss_as_Z1":      pct(z1 + M_total, n + M_total),
        "diff_baseline_vs_nominal_pp": round(pct(z1, n) - expected_z1*100, 2),
        "identifying":          ident,
    })

res = pd.DataFrame(rows)
print("\n=== realized vs nominal-under-our-Z-mapping, per phase ===")
print(res.to_string(index=False))
res.to_csv(ROOT/"analysis/diagnose_missing_mentors/rollout_realized_with_bounds.tsv",
            sep="\t", index=False)
print("\nsaved -> rollout_realized_with_bounds.tsv")

# =============================================================================
# Audit group B: users WITH a mentor row in snapshot but NO local enwiki reg_ts.
# Expected: 927 (= 7,378 reg_ts MISSING - 6,451 no_row deleted).
# =============================================================================
in_snap = ds["userId"].map(lambda u: u in snap)
groupB  = ds[in_snap & ds["reg_ts"].isna()].copy()
print(f"group B size (has_row & reg_ts NaN): {len(groupB):,}")

print("\n--- by mentorshipState ---")
print(groupB["mentorshipState"].value_counts(dropna=False).to_string())

print("\n--- by Z (unset|50 vs 0) ---")
print(groupB["mentorshipState"].isin(["unset","50"]).value_counts().rename({True:"Z=1",False:"Z=0"}).to_string())

def assigned_ts(uid):
    v = snap.get(uid)
    if v is None: return None
    for x in (v if isinstance(v,(list,tuple)) else [v]):
        if isinstance(x,str) and len(x)>=8 and x[:4].isdigit():
            return x
    return None
groupB["assigned_ts"]    = groupB["userId"].map(assigned_ts)
groupB["assigned_month"] = pd.to_datetime(groupB["assigned_ts"], errors="coerce").dt.to_period("M").astype("string")

print("\n--- by assigned_month ---")
print(groupB["assigned_month"].value_counts(dropna=False).sort_index().to_string())

OUT = ROOT / "analysis/diagnose_missing_mentors/group_B_has_row_no_reg.tsv"
groupB[["userId","mentorshipState","assigned_ts"]].to_csv(OUT, sep="\t", index=False)
print(f"\nwritten -> {OUT}")

n_missing = int(ds["reg_ts"].isna().sum())
n_groupA  = int((~in_snap & ds["reg_ts"].isna()).sum())
print(f"\nsanity: reg_ts NaN total={n_missing:,} = A(no_row & NaN)={n_groupA:,} + B(in_snap & NaN)={len(groupB):,}")