#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
regression-to-the-mean.py
Science Journaling Club, Volume 2, Issue 2, Winter 2026, "Statistics That Fool You".

QUESTION
--------
Select the worst performers on a first measurement, do nothing whatsoever to them,
measure them again, and they improve. How much of a typical before-and-after
improvement can regression to the mean explain on its own, with no treatment effect
at all? And when a genuine treatment effect IS present, how badly does a single-arm
before-and-after design overstate it compared with a randomised design run on the
same simulated people?

MODEL
-----
Every simulated individual carries a stable underlying trait T and is measured twice.
Each measurement adds fresh, independent noise:

    T   ~ Normal(0, sigma_T^2)          the trait, fixed for that individual
    X1  = T + e1,   e1 ~ Normal(0, sigma_e^2)
    X2  = T + e2,   e2 ~ Normal(0, sigma_e^2)      e1, e2 independent of each other

The noise-to-signal ratio is lam = sigma_e / sigma_T. The correlation between the two
measurements, which in the measurement literature is the test-retest reliability, is

    r = sigma_T^2 / (sigma_T^2 + sigma_e^2) = 1 / (1 + lam^2)

We fix the total variance of a single measurement at 1 by setting
sigma_T = sqrt(r) and sigma_e = sqrt(1 - r), so every quantity below is already in
units of the population standard deviation of one measurement. That is the unit a
practitioner actually sees: the spread of the whole group on a single testing day.

Selection takes the bottom fraction p of the population on X1 (the "worst cases").
No treatment is applied. We then report

    Delta = mean(X2 - X1) over the selected group

which is the apparent improvement, in population SDs, produced by nothing at all.

CLOSED FORM USED FOR VALIDATION
-------------------------------
(X1, X2) is bivariate normal, both marginals standard, correlation r. Therefore
E[X2 | X1] = r * X1 exactly, and so for ANY selection rule based on X1 alone,

    E[X2 | selected] = r * E[X1 | selected]                          (shrinkage law)

The group's mean is pulled toward the population mean by exactly the factor r. For a
bottom-p cut at z_p = Phi^{-1}(p), the standard truncated-normal mean gives

    E[X1 | selected] = -phi(z_p) / p
    E[X2 | selected] = -r * phi(z_p) / p
    Delta_analytic   = (1 - r) * phi(z_p) / p

With zero noise, lam = 0 so r = 1, and Delta_analytic is exactly 0: X2 = X1 for
everybody and no spurious improvement can appear. Both the shrinkage factor and
Delta are printed beside the simulation for every cell of the grid.

TREATMENT PART
--------------
A true treatment effect tau (in population SDs) is added to the second measurement of
anyone who receives treatment. Two designs are run on the same simulated people:

  single-arm before-and-after : everybody selected is treated; estimate = mean(X2-X1)
                                expected value = tau + (1-r)*phi(z_p)/p
  randomised controlled       : the selected group is split at random into treated and
                                control; estimate = mean(X2 | treated) - mean(X2 | control)
                                expected value = tau, because both arms regress equally

ASSUMPTIONS
-----------
 * The trait is perfectly stable between the two measurements. No real change, no
   ageing, no learning, no seasonal drift, no genuine natural history of a disease.
 * Noise is normal, additive, homoscedastic, equal on both occasions, and independent
   between occasions.
 * Selection uses the first measurement alone and nothing else.
 * The treatment effect is a constant shift, the same for everyone, with no
   interaction with the trait and no effect on measurement noise.
 * No dropout, no missing data, no measurement floor or ceiling.

LIMITATIONS
-----------
 * Real second measurements are rarely pure repeats. Anything that genuinely changes
   between them (recovery, growth, practice, a real seasonal cycle) adds to Delta and
   is not separated from it here.
 * Normality matters. A heavy-tailed noise distribution changes the truncated-mean
   factor phi(z_p)/p, and the shrinkage law E[X2|X1] = r*X1 is a property of the
   bivariate normal, not of every distribution with correlation r.
 * Reliability r is treated as known. In practice it must be estimated, and a
   regression-to-the-mean correction is only as good as that estimate.
 * The model says nothing about how to choose r for any real instrument. It produces
   the size of the artefact given r; it cannot tell you what r your clinic's blood
   pressure cuff or your school's quiz has.

RUN
---
    python regression-to-the-mean.py > regression-to-the-mean-output.txt

Python 3.12, numpy. Master seed printed below and hard-coded at MASTER_SEED.
"""

import math
import sys
import time
from statistics import NormalDist

import numpy as np

# --------------------------------------------------------------------------- #
# Configuration
# --------------------------------------------------------------------------- #

MASTER_SEED = 20260211

# Selection thresholds: the bottom fraction taken on the first measurement.
P_GRID = [0.50, 0.25, 0.10, 0.05, 0.02, 0.01]

# Noise-to-signal ratios sigma_e / sigma_T. lam = 0 is the zero-noise control.
LAM_GRID = [0.0, 0.25, 0.50, 0.75, 1.00, 1.50, 2.00]

N_PER_REP = 100_000      # individuals simulated per replicate
N_REPS = 60              # independent replicates per grid cell

# A cell whose |z| exceeds this is re-run on a completely fresh stream before we
# are willing to call it a disagreement. Set once, before looking at any output.
RERUN_Z = 3.0

# The cell the article uses as its headline: bottom 10% on a measurement whose
# reliability is 0.5 (noise variance equal to trait variance, lam = 1).
HEAD_P = 0.10
HEAD_LAM = 1.00

# Treatment-effect grid, in population SDs.
TAU_GRID = [0.0, 0.10, 0.20, 0.30, 0.50, 1.00]

# Scenarios for the trial comparison: (label, p, lam)
TRIAL_SCENARIOS = [
    ("tight cut, clean measure", 0.05, 0.50),
    ("headline cell", 0.10, 1.00),
    ("loose cut, noisy measure", 0.25, 1.50),
]

ND = NormalDist()


def phi(x):
    """Standard normal density."""
    return math.exp(-0.5 * x * x) / math.sqrt(2.0 * math.pi)


def z_of(p):
    """Standard normal quantile."""
    return ND.inv_cdf(p)


def rule(ch="-", n=78):
    return ch * n


def head(title):
    print()
    print(rule("="))
    print(title)
    print(rule("="))


# --------------------------------------------------------------------------- #
# Core simulation of one grid cell
# --------------------------------------------------------------------------- #

def run_cell(p, lam, seed_seq, n_per_rep=N_PER_REP, n_reps=N_REPS,
             keep_running=False, scatter_n=0):
    """
    Simulate n_reps independent replicates of n_per_rep individuals.

    Returns a dict of replicate-level estimates. Standard errors come from the
    spread across replicates (a batched-means estimator), so they are measured
    from the runs rather than assumed.
    """
    r = 1.0 / (1.0 + lam * lam)
    sd_t = math.sqrt(r)
    sd_e = math.sqrt(1.0 - r)

    rng = np.random.Generator(np.random.PCG64(seed_seq))

    m1 = np.empty(n_reps)     # mean X1 of the selected group
    m2 = np.empty(n_reps)     # mean X2 of the selected group
    dl = np.empty(n_reps)     # mean (X2 - X1) of the selected group
    sh = np.empty(n_reps)     # measured shrinkage, m2 / m1
    rr = np.empty(n_reps)     # measured Pearson correlation of X1 and X2
    sd_pre = np.empty(n_reps)  # SD of X1 inside the selected group
    nsel = np.empty(n_reps, dtype=np.int64)

    running_sum_d = 0.0
    running_n = 0
    running = []              # (cumulative individuals selected, running Delta)

    scatter = None

    k = max(1, int(round(p * n_per_rep)))

    for i in range(n_reps):
        t = rng.standard_normal(n_per_rep) * sd_t
        x1 = t + rng.standard_normal(n_per_rep) * sd_e
        x2 = t + rng.standard_normal(n_per_rep) * sd_e

        # Bottom k on the first measurement. argpartition avoids a full sort.
        idx = np.argpartition(x1, k - 1)[:k]
        s1 = x1[idx]
        s2 = x2[idx]

        m1[i] = s1.mean()
        m2[i] = s2.mean()
        dl[i] = (s2 - s1).mean()
        sh[i] = m2[i] / m1[i]
        sd_pre[i] = s1.std(ddof=1)
        nsel[i] = k
        rr[i] = np.corrcoef(x1, x2)[0, 1]

        if keep_running:
            running_sum_d += float((s2 - s1).sum())
            running_n += k
            running.append((running_n, running_sum_d / running_n))

        if scatter_n and scatter is None:
            take = rng.choice(n_per_rep, size=scatter_n, replace=False)
            scatter = (x1[take].copy(), x2[take].copy())

    def mse(a):
        """Mean and standard error of the mean, both from the replicate spread."""
        return float(a.mean()), float(a.std(ddof=1) / math.sqrt(len(a)))

    out = {
        "p": p, "lam": lam, "r": r, "k": k,
        "n_total": n_per_rep * n_reps,
        "n_sel_total": int(nsel.sum()),
    }
    out["m1"], out["m1_se"] = mse(m1)
    out["m2"], out["m2_se"] = mse(m2)
    out["delta"], out["delta_se"] = mse(dl)
    out["shrink"], out["shrink_se"] = mse(sh)
    out["corr"], out["corr_se"] = mse(rr)
    out["sd_pre"], out["sd_pre_se"] = mse(sd_pre)
    out["running"] = running
    out["scatter"] = scatter
    out["reps_delta"] = dl.copy()
    return out


def analytic(p, lam):
    r = 1.0 / (1.0 + lam * lam)
    zp = z_of(p)
    a = phi(zp) / p                     # = -E[X1 | selected]
    m1 = -a
    m2 = -r * a
    delta = (1.0 - r) * a
    # SD of X1 inside the truncated group. For X ~ N(0,1) cut at X < z_p,
    # Var = 1 - z_p*phi(z_p)/p - (phi(z_p)/p)^2, and z_p is negative here.
    var_t = 1.0 - zp * a - a * a
    return {"r": r, "zp": zp, "m1": m1, "m2": m2, "delta": delta,
            "sd_pre": math.sqrt(max(var_t, 0.0))}


def zscore(obs, exp, se):
    if se <= 0:
        return float("nan")
    return (obs - exp) / se


def chi2_sf(x, k):
    """Upper tail of the chi-square distribution, k degrees of freedom.

    Wilson-Hilferty cube-root transform, which is accurate to better than 1e-3
    in the tail for k >= 30 and is all we need to report a goodness-of-fit p.
    """
    if k <= 0:
        return float("nan")
    t = (x / k) ** (1.0 / 3.0)
    m = 1.0 - 2.0 / (9.0 * k)
    s = math.sqrt(2.0 / (9.0 * k))
    return 1.0 - ND.cdf((t - m) / s)


def t_crit_two_sided(alpha, df):
    """Two-sided critical value of Student t, by bisection on the CDF.

    Uses the incomplete-beta-free route: t's CDF is obtained by numerically
    integrating its density with Simpson's rule, which is plenty for a printed
    threshold and keeps the dependency list at numpy alone.
    """
    def tcdf(x):
        n = 4000
        lo, hi = -40.0, x
        h = (hi - lo) / n
        c = (math.lgamma((df + 1) / 2.0) - math.lgamma(df / 2.0)
             - 0.5 * math.log(df * math.pi))
        def f(u):
            return math.exp(c - (df + 1) / 2.0 * math.log1p(u * u / df))
        tot = f(lo) + f(hi)
        for i in range(1, n):
            tot += (4 if i % 2 else 2) * f(lo + i * h)
        return tot * h / 3.0
    lo, hi = 0.0, 60.0
    for _ in range(200):
        mid = 0.5 * (lo + hi)
        if 2.0 * (1.0 - tcdf(mid)) > alpha:
            lo = mid
        else:
            hi = mid
    return 0.5 * (lo + hi)


# --------------------------------------------------------------------------- #
# Trial comparison: single-arm before-and-after against a randomised design
# --------------------------------------------------------------------------- #

def run_trial(p, lam, tau, seed_seq, n_per_rep=N_PER_REP, n_reps=N_REPS):
    """
    Same population, same selection, then two designs run side by side.

    single-arm : every selected person is treated; estimate = mean(X2 - X1)
    randomised : the selected group is split 50/50 at random; treated get +tau on
                 the second measurement; estimate = mean(X2|T) - mean(X2|C).
                 Also reported: the change-score (difference-in-differences) form.
    """
    r = 1.0 / (1.0 + lam * lam)
    sd_t = math.sqrt(r)
    sd_e = math.sqrt(1.0 - r)
    rng = np.random.Generator(np.random.PCG64(seed_seq))

    single = np.empty(n_reps)
    rct = np.empty(n_reps)
    did = np.empty(n_reps)
    sig = np.zeros(n_reps)          # single-arm paired t reaches p < 0.05?
    sig_rct = np.zeros(n_reps)

    k = max(2, int(round(p * n_per_rep)))
    if k % 2:
        k -= 1

    for i in range(n_reps):
        t = rng.standard_normal(n_per_rep) * sd_t
        x1 = t + rng.standard_normal(n_per_rep) * sd_e
        x2 = t + rng.standard_normal(n_per_rep) * sd_e

        idx = np.argpartition(x1, k - 1)[:k]
        s1 = x1[idx]
        s2 = x2[idx]

        # --- single-arm: everybody treated ---
        d_all = (s2 + tau) - s1
        single[i] = d_all.mean()
        se_d = d_all.std(ddof=1) / math.sqrt(k)
        sig[i] = 1.0 if abs(d_all.mean() / se_d) > 1.96 else 0.0

        # --- randomised: half treated, half not ---
        perm = rng.permutation(k)
        tre = perm[: k // 2]
        con = perm[k // 2:]
        y_t = s2[tre] + tau
        y_c = s2[con]
        rct[i] = y_t.mean() - y_c.mean()
        se_r = math.sqrt(y_t.var(ddof=1) / len(y_t) + y_c.var(ddof=1) / len(y_c))
        sig_rct[i] = 1.0 if abs(rct[i] / se_r) > 1.96 else 0.0

        d_t = y_t - s1[tre]
        d_c = y_c - s1[con]
        did[i] = d_t.mean() - d_c.mean()

    def mse(a):
        return float(a.mean()), float(a.std(ddof=1) / math.sqrt(len(a)))

    out = {"p": p, "lam": lam, "r": r, "tau": tau, "k": k}
    out["single"], out["single_se"] = mse(single)
    out["rct"], out["rct_se"] = mse(rct)
    out["did"], out["did_se"] = mse(did)
    out["sig_rate"] = float(sig.mean())
    out["sig_rate_rct"] = float(sig_rct.mean())
    return out


# --------------------------------------------------------------------------- #
# Main
# --------------------------------------------------------------------------- #

def main():
    t_start = time.time()
    ss = np.random.SeedSequence(MASTER_SEED)
    # One independent child stream per job, spawned in a fixed order.
    n_cells = len(P_GRID) * len(LAM_GRID)
    n_trials = len(TRIAL_SCENARIOS) * len(TAU_GRID)
    children = ss.spawn(n_cells + n_trials + 8)
    ci = 0

    print(rule("="))
    print("REGRESSION TO THE MEAN: HOW MUCH IMPROVEMENT COMES FROM NOTHING")
    print("Science Journaling Club, Volume 2 Issue 2, Winter 2026")
    print(rule("="))
    print(f"master seed      : {MASTER_SEED}")
    print(f"generator        : numpy PCG64 via SeedSequence.spawn, one stream per cell")
    print(f"python           : {sys.version.split()[0]}")
    print(f"numpy            : {np.__version__}")
    print(f"individuals/rep  : {N_PER_REP:,}")
    print(f"replicates/cell  : {N_REPS}")
    print(f"individuals/cell : {N_PER_REP * N_REPS:,}")
    print(f"grid             : {len(P_GRID)} selection thresholds x {len(LAM_GRID)} "
          f"noise ratios = {n_cells} cells")
    print()
    print("Model: T ~ N(0, r), e ~ N(0, 1-r), X1 = T + e1, X2 = T + e2.")
    print("       Total variance of one measurement is 1, so every number below is in")
    print("       units of the population SD of a single measurement.")
    print("       r = 1/(1+lam^2) is the correlation between the two measurements.")
    print("       Selection takes the bottom p on X1. NO TREATMENT IS APPLIED in Part A.")
    print()
    print("All standard errors are computed from the spread across the 25 independent")
    print("replicates of each cell (batched means), not from an assumed formula.")

    # ----------------------------------------------------------------- Part A
    head("PART A.1  VALIDATION OF THE SHRINKAGE LAW   E[X2|sel] = r * E[X1|sel]")
    print("For a bivariate normal with correlation r, the mean of a group selected on")
    print("the first measurement is pulled toward the population mean by exactly r,")
    print("whatever the selection rule, as long as it uses X1 alone.")
    print()
    print(f"{'p':>6} {'lam':>5} {'r (exact)':>10} {'shrink sim':>12} {'SE':>9} "
          f"{'difference':>11} {'z':>7} {'corr(X1,X2)':>12}")
    print(rule())

    cells = {}
    max_abs_z_shrink = 0.0
    worst_shrink = None
    for p in P_GRID:
        for lam in LAM_GRID:
            c = run_cell(p, lam, children[ci]); ci += 1
            a = analytic(p, lam)
            cells[(p, lam)] = (c, a)
            z = zscore(c["shrink"], a["r"], c["shrink_se"])
            if math.isfinite(z) and abs(z) > max_abs_z_shrink:
                max_abs_z_shrink = abs(z)
                worst_shrink = (p, lam, c["shrink"], a["r"], c["shrink_se"], z)
            zs = "  exact" if c["shrink_se"] == 0 else f"{z:7.2f}"
            print(f"{p:6.2f} {lam:5.2f} {a['r']:10.6f} {c['shrink']:12.6f} "
                  f"{c['shrink_se']:9.6f} {c['shrink'] - a['r']:+11.6f} {zs} "
                  f"{c['corr']:12.6f}")
        print(rule("."))

    print()
    print(f"Largest |z| on the shrinkage factor across all {n_cells} cells: "
          f"{max_abs_z_shrink:.2f}")
    if worst_shrink:
        wp, wl, wo, we, ws, wz = worst_shrink
        print(f"  worst cell p={wp}, lam={wl}: simulated {wo:.6f} vs exact {we:.6f}, "
              f"SE {ws:.6f}, z = {wz:+.2f}")
    tc = t_crit_two_sided(0.05 / n_cells, N_REPS - 1)
    print(f"Bonferroni threshold for {n_cells} cells at family-wise 5%, Student t with "
          f"{N_REPS - 1} df: {tc:.2f}")

    head("PART A.2  THE SPURIOUS IMPROVEMENT, NO TREATMENT AT ALL")
    print("Delta = mean(X2 - X1) over the selected group, in population SDs.")
    print("Analytic value: Delta = (1 - r) * phi(z_p) / p.")
    print()
    print(f"{'p':>6} {'lam':>5} {'r':>8} {'pre-mean':>10} {'post-mean':>10} "
          f"{'Delta sim':>10} {'SE':>9} {'Delta exact':>12} {'diff':>10} {'z':>7} {'d_within':>9}")
    print(rule("-", 116))

    max_abs_z_delta = 0.0
    worst_delta = None
    for p in P_GRID:
        for lam in LAM_GRID:
            c, a = cells[(p, lam)]
            z = zscore(c["delta"], a["delta"], c["delta_se"])
            if math.isfinite(z) and abs(z) > max_abs_z_delta:
                max_abs_z_delta = abs(z)
                worst_delta = (p, lam, c["delta"], a["delta"], c["delta_se"], z)
            zs = "  exact" if c["delta_se"] == 0 else f"{z:7.2f}"
            d_within = c["delta"] / c["sd_pre"] if c["sd_pre"] > 0 else float("nan")
            dw = f"{d_within:9.4f}" if math.isfinite(d_within) else "      n/a"
            print(f"{p:6.2f} {lam:5.2f} {a['r']:8.4f} {c['m1']:10.4f} {c['m2']:10.4f} "
                  f"{c['delta']:10.5f} {c['delta_se']:9.5f} {a['delta']:12.5f} "
                  f"{c['delta'] - a['delta']:+10.5f} {zs} {dw}")
        print(rule("."))

    print()
    print("d_within is the spurious improvement divided by the SD of the selected")
    print("group's own first measurement, which is the effect size a single-arm paper")
    print("would usually report.")
    print()
    print(f"Largest |z| on Delta across all {n_cells} cells: {max_abs_z_delta:.2f}")
    if worst_delta:
        wp, wl, wo, we, ws, wz = worst_delta
        print(f"  worst cell p={wp}, lam={wl}: simulated {wo:.5f} vs exact {we:.5f}, "
              f"SE {ws:.5f}, z = {wz:+.2f}")
    print(f"  Bonferroni threshold for {n_cells} cells at family-wise 5%: {tc:.2f}")

    # Pooled goodness of fit. The 36 noisy cells only; the six zero-noise cells
    # have SE identically zero and contribute no z.
    zs_all = []
    for p in P_GRID:
        for lam in LAM_GRID:
            c, a = cells[(p, lam)]
            zz = zscore(c["delta"], a["delta"], c["delta_se"])
            if math.isfinite(zz):
                zs_all.append(zz)
    chi = sum(z * z for z in zs_all)
    dfree = len(zs_all)
    print()
    print(f"Pooled goodness of fit over the {dfree} cells that carry a standard error:")
    print(f"  sum of z^2 = {chi:.2f} on {dfree} degrees of freedom, "
          f"p = {chi2_sf(chi, dfree):.3f}")
    print(f"  mean z = {sum(zs_all) / dfree:+.3f}, "
          f"SD of z = {math.sqrt(sum((z - sum(zs_all) / dfree) ** 2 for z in zs_all) / (dfree - 1)):.3f}")
    print("  (a correct implementation should give sum z^2 near the degrees of")
    print("   freedom, mean z near 0 and SD of z near 1)")

    # --- pre-registered re-run of any cell past RERUN_Z, on a fresh stream ---
    print()
    print(f"RE-RUN RULE, fixed before the run: any cell with |z| > {RERUN_Z:.1f} on Delta")
    print("is simulated again from a completely independent stream. A real disagreement")
    print("repeats; a fluctuation does not.")
    flagged = []
    for p in P_GRID:
        for lam in LAM_GRID:
            c, a = cells[(p, lam)]
            zz = zscore(c["delta"], a["delta"], c["delta_se"])
            if math.isfinite(zz) and abs(zz) > RERUN_Z:
                flagged.append((p, lam, zz))
    if not flagged:
        print(f"  no cell exceeded |z| = {RERUN_Z:.1f}. Nothing to re-run.")
    else:
        rerun_ss = np.random.SeedSequence(MASTER_SEED + 777).spawn(len(flagged))
        print(f"  {len(flagged)} cell(s) flagged.")
        print(f"{'p':>6} {'lam':>5} {'z first run':>12} {'Delta re-run':>13} "
              f"{'SE':>9} {'exact':>10} {'z re-run':>9} {'verdict':>12}")
        for j, (p, lam, zz) in enumerate(flagged):
            c2 = run_cell(p, lam, rerun_ss[j])
            a = analytic(p, lam)
            z2 = zscore(c2["delta"], a["delta"], c2["delta_se"])
            verdict = "repeats" if abs(z2) > RERUN_Z and z2 * zz > 0 else "fluctuation"
            print(f"{p:6.2f} {lam:5.2f} {zz:12.2f} {c2['delta']:13.5f} "
                  f"{c2['delta_se']:9.5f} {a['delta']:10.5f} {z2:9.2f} {verdict:>12}")

    head("PART A.3  THE ZERO-NOISE CONTROL  (lam = 0, r = 1)")
    print("With no measurement noise the second measurement equals the first exactly,")
    print("so no spurious improvement can appear. Every p is checked.")
    print()
    print(f"{'p':>6} {'r':>6} {'Delta sim':>14} {'SE':>10} {'Delta exact':>13} {'verdict':>12}")
    print(rule())
    zero_ok = True
    for p in P_GRID:
        c, a = cells[(p, 0.0)]
        ok = (c["delta"] == 0.0 and c["delta_se"] == 0.0)
        zero_ok = zero_ok and ok
        print(f"{p:6.2f} {a['r']:6.1f} {c['delta']:14.10f} {c['delta_se']:10.6f} "
              f"{a['delta']:13.10f} {'exactly 0' if ok else 'NONZERO':>12}")
    print()
    print("zero-noise control: " + ("PASSED, all Delta identically zero"
                                    if zero_ok else "FAILED"))

    # headline cell detail
    hc, ha = cells[(HEAD_P, HEAD_LAM)]
    head("PART A.4  THE HEADLINE CELL IN FULL")
    print(f"Selection: bottom {HEAD_P:.0%} on the first measurement.")
    print(f"Noise-to-signal ratio lam = {HEAD_LAM:.2f}, so reliability r = {ha['r']:.4f}.")
    print(f"Individuals simulated: {hc['n_total']:,} in {N_REPS} replicates.")
    print(f"Selected per replicate: {hc['k']:,}   total selected: {hc['n_sel_total']:,}")
    print()
    print(f"  mean first measurement, selected  : {hc['m1']:+.4f} "
          f"+/- {hc['m1_se']:.4f}   (exact {ha['m1']:+.4f})")
    print(f"  mean second measurement, selected : {hc['m2']:+.4f} "
          f"+/- {hc['m2_se']:.4f}   (exact {ha['m2']:+.4f})")
    print(f"  apparent improvement Delta        : {hc['delta']:+.4f} "
          f"+/- {hc['delta_se']:.4f}   (exact {ha['delta']:+.4f})")
    print(f"  shrinkage m2/m1                   : {hc['shrink']:.6f} "
          f"+/- {hc['shrink_se']:.6f}   (exact r {ha['r']:.6f})")
    print(f"  SD of selected group on X1        : {hc['sd_pre']:.4f}   "
          f"(exact {ha['sd_pre']:.4f})")
    print(f"  improvement / that SD             : {hc['delta'] / hc['sd_pre']:.4f}")
    print()
    print(f"  percent of the selected group's gap to the mean closed by doing nothing:")
    print(f"    {100.0 * (1.0 - ha['r']):.1f}%  (exactly 1 - r)")

    # ----------------------------------------------------------------- Part B
    head("PART B  A REAL TREATMENT EFFECT, TWO DESIGNS, SAME SIMULATED PEOPLE")
    print("tau is the true effect in population SDs, added to the second measurement")
    print("of anyone treated. Single-arm treats everyone selected and reports")
    print("mean(X2 - X1). Randomised splits the selected group 50/50 and reports")
    print("mean(X2 | treated) - mean(X2 | control).")
    print()

    trials = {}
    max_abs_z_rct = 0.0
    worst_rct = None
    for label, p, lam in TRIAL_SCENARIOS:
        a = analytic(p, lam)
        print(rule("-", 110))
        print(f"SCENARIO: {label}   p = {p:.2f}, lam = {lam:.2f}, r = {a['r']:.4f}, "
              f"regression bias = {a['delta']:.4f} SD")
        print(rule("-", 110))
        print(f"{'tau':>6} {'single-arm':>11} {'SE':>8} {'expected':>10} "
              f"{'bias':>8} {'overstate':>10} {'RCT est':>9} {'SE':>8} {'z(RCT)':>7} "
              f"{'DiD est':>9} {'sig%':>6}")
        for tau in TAU_GRID:
            tr = run_trial(p, lam, tau, children[ci]); ci += 1
            trials[(p, lam, tau)] = tr
            exp_single = tau + a["delta"]
            over = (tr["single"] / tau) if tau > 0 else float("inf")
            zr = zscore(tr["rct"], tau, tr["rct_se"])
            if math.isfinite(zr) and abs(zr) > max_abs_z_rct:
                max_abs_z_rct = abs(zr)
                worst_rct = (p, lam, tau, tr["rct"], tau, tr["rct_se"], zr)
            over_s = "     inf" if tau == 0 else f"{over:9.2f}x"
            print(f"{tau:6.2f} {tr['single']:11.5f} {tr['single_se']:8.5f} "
                  f"{exp_single:10.5f} {a['delta']:8.4f} {over_s:>10} "
                  f"{tr['rct']:9.5f} {tr['rct_se']:8.5f} {zr:7.2f} "
                  f"{tr['did']:9.5f} {100 * tr['sig_rate']:5.0f}%")
        print()

    print(f"Largest |z| on the randomised estimate against the true tau, across all "
          f"{n_trials} trial cells: {max_abs_z_rct:.2f}")
    if worst_rct:
        wp, wl, wt, wo, we, ws, wz = worst_rct
        print(f"  worst: p={wp}, lam={wl}, tau={wt}: RCT {wo:.5f} vs true {we:.5f}, "
              f"SE {ws:.5f}, z = {wz:+.2f}")
    print()
    print("sig% is the share of replicates in which the single-arm paired t statistic")
    print("exceeded 1.96 in absolute value. At tau = 0 it is the false-positive rate of")
    print("a before-and-after design applied to a selected group, and it is not 5%.")

    head("PART B.2  THE SHARE OF THE APPARENT EFFECT THAT IS ARTEFACT")
    print(f"{'scenario':>28} {'r':>7} {'tau':>6} {'single-arm':>11} "
          f"{'artefact share':>15} {'inflation':>10}")
    print(rule("-", 84))
    for label, p, lam in TRIAL_SCENARIOS:
        a = analytic(p, lam)
        for tau in TAU_GRID:
            if tau == 0:
                continue
            tr = trials[(p, lam, tau)]
            share = a["delta"] / tr["single"]
            print(f"{label:>28} {a['r']:7.4f} {tau:6.2f} {tr['single']:11.5f} "
                  f"{100 * share:14.1f}% {tr['single'] / tau:9.2f}x")
        print(rule("."))

    # --------------------------------------------------------- Club homework
    head("PART C  THE CLUB'S OWN HOMEWORK, SIMULATED")
    print("A quiz taken twice by a class of 30. This is a SIMULATION with class-sized")
    print("numbers, not our actual marks; the club has no such dataset and did not")
    print("collect one. Reliability of a short classroom quiz is taken as r = 0.70,")
    print("which is a typical published value for a short classroom test. The bottom")
    print("8 of 30 are picked out for 'extra help' that in this model does nothing.")
    class_n = 30
    bottom_k = 8
    r_quiz = 0.70
    n_classes = 200_000
    rngc = np.random.Generator(np.random.PCG64(children[ci])); ci += 1
    sd_t = math.sqrt(r_quiz)
    sd_e = math.sqrt(1.0 - r_quiz)
    t = rngc.standard_normal((n_classes, class_n)) * sd_t
    q1 = t + rngc.standard_normal((n_classes, class_n)) * sd_e
    q2 = t + rngc.standard_normal((n_classes, class_n)) * sd_e
    order = np.argsort(q1, axis=1)[:, :bottom_k]
    rows = np.arange(n_classes)[:, None]
    g1 = q1[rows, order]
    g2 = q2[rows, order]
    gain = (g2 - g1).mean(axis=1)
    # The whole class, and the 22 who were not selected. The unselected group is
    # the honest comparison, and it moves the other way.
    all_gain = (q2 - q1).mean(axis=1)
    rest_mask = np.ones((n_classes, class_n), dtype=bool)
    rest_mask[rows, order] = False
    rest_gain = ((q2 - q1) * rest_mask).sum(axis=1) / (class_n - bottom_k)
    top_order = np.argsort(q1, axis=1)[:, -bottom_k:]
    t1 = q1[rows, top_order]
    t2 = q2[rows, top_order]
    top_gain = (t2 - t1).mean(axis=1)
    # Put it on a 100-point scale with SD 12, a plausible classroom spread
    SCALE_SD = 12.0
    SCALE_MEAN = 68.0
    print()
    print(f"  classes simulated            : {n_classes:,}")
    print(f"  class size                   : {class_n}, bottom {bottom_k} selected "
          f"(p = {bottom_k / class_n:.4f})")
    print(f"  quiz reliability r           : {r_quiz:.2f}")
    zq = z_of(bottom_k / class_n)
    delta_q = (1 - r_quiz) * phi(zq) / (bottom_k / class_n)
    print(f"  mean gain of the bottom 8    : {gain.mean():+.4f} SD "
          f"+/- {gain.std(ddof=1) / math.sqrt(n_classes):.5f}")
    print(f"  continuous-population exact  : {delta_q:+.4f} SD "
          f"(a finite class of 30 differs slightly; order statistics, not a fixed cut)")
    print(f"  mean gain of the whole class : {all_gain.mean():+.6f} SD "
          f"+/- {all_gain.std(ddof=1) / math.sqrt(n_classes):.6f}")
    print(f"  mean gain of the other 22    : {rest_gain.mean():+.4f} SD "
          f"+/- {rest_gain.std(ddof=1) / math.sqrt(n_classes):.5f}")
    print(f"  mean gain of the TOP 8       : {top_gain.mean():+.4f} SD "
          f"+/- {top_gain.std(ddof=1) / math.sqrt(n_classes):.5f}   "
          f"(the same effect, running downhill)")
    print(f"  on a 100-point scale, mean {SCALE_MEAN:.0f} and SD {SCALE_SD:.0f}:")
    print(f"    bottom 8 first attempt     : {SCALE_MEAN + SCALE_SD * g1.mean():.1f} marks")
    print(f"    bottom 8 second attempt    : {SCALE_MEAN + SCALE_SD * g2.mean():.1f} marks")
    print(f"    apparent improvement       : {SCALE_SD * gain.mean():+.1f} marks")
    print(f"    whole class moved          : {SCALE_SD * all_gain.mean():+.2f} marks")
    print(f"    the other 22 moved         : {SCALE_SD * rest_gain.mean():+.2f} marks")
    print(f"    the top 8 moved            : {SCALE_SD * top_gain.mean():+.2f} marks")
    print(f"  share of classes in which the bottom 8 improved : "
          f"{100 * (gain > 0).mean():.1f}%")
    print(f"  share of classes in which the top 8 got worse   : "
          f"{100 * (top_gain < 0).mean():.1f}%")
    print()
    print("  Both halves of that are the same arithmetic. Help the bottom 8 and they")
    print("  improve. Leave the top 8 alone and they decline. Neither group was touched.")

    # ------------------------------------------------------- Literature check
    head("PART E  SET AGAINST A PUBLISHED NUMBER")
    print("Krogsboll, Hrobjartsson and Gotzsche (2009) pooled 37 three-armed trials in")
    print("2,900 patients across 8 clinical conditions, in which one arm received no")
    print("treatment at all. Change from baseline in those untreated arms had a pooled")
    print("standardised mean difference of -0.24. Their active-treatment arms gave")
    print("-1.01 and their placebo arms -0.44, so the untreated arms accounted for")
    print("about 24% of the movement seen under active treatment.")
    print()
    print("Their -0.24 is real patients and mixes several causes: natural history of")
    print("the illness, regression to the mean, and anything else that changes between")
    print("two visits. Our model produces only the regression part. So the fair")
    print("question is: what reliability would a trial's entry measurement need for")
    print("regression alone to account for the whole -0.24?")
    print()
    OBS = 0.24
    print(f"{'entry cut p':>12} {'phi(z_p)/p':>12} {'r needed':>10} "
          f"{'lam needed':>11} {'our Delta at r=0.5':>20} {'our Delta at r=0.8':>20}")
    print(rule("-", 92))
    for p in P_GRID:
        f_p = phi(z_of(p)) / p
        r_need = 1.0 - OBS / f_p
        if r_need <= 0:
            r_s, l_s = "  <= 0", "     n/a"
        else:
            lam_need = math.sqrt(1.0 / r_need - 1.0)
            r_s, l_s = f"{r_need:10.4f}", f"{lam_need:11.4f}"
        print(f"{p:12.2f} {f_p:12.4f} {r_s:>10} {l_s:>11} "
              f"{0.5 * f_p:20.4f} {0.2 * f_p:20.4f}")
    print()
    print("Read the table this way. If a trial enrols the worst 10% on a single entry")
    print("measurement, regression alone matches the published -0.24 as soon as that")
    r10 = 1.0 - OBS / (phi(z_of(0.10)) / 0.10)
    print(f"measurement's test-retest reliability is {r10:.3f} or worse. Reliabilities")
    print("below that are ordinary for single-occasion clinical and behavioural")
    print("measurements. The published number therefore sets no lower bound on how much")
    print("of it is natural history: regression can supply all of it on plausible")
    print("numbers, and our model cannot tell you how much it actually did supply in")
    print("any particular trial. What it can say is that a single-arm design has no way")
    print("of separating the two, and a randomised design does not need to.")
    print()
    print("The opposite reading matters as much. If entry is loose, say the worst half,")
    f50 = phi(z_of(0.50)) / 0.50
    print(f"then phi(z_p)/p is only {f50:.4f} and regression can reach -0.24 only if")
    print(f"reliability falls to {1.0 - OBS / f50:.3f}. Tight entry criteria are what make")
    print("the artefact large. That is a design choice, not a property of the disease.")

    # ------------------------------------------------------------ Convergence
    head("PART D  CONVERGENCE OF THE HEADLINE CELL")
    print("The running estimate of Delta as replicates accumulate, for the headline")
    print(f"cell p = {HEAD_P:.2f}, lam = {HEAD_LAM:.2f}. Each replicate adds "
          f"{int(round(HEAD_P * N_PER_REP)):,} selected individuals.")
    conv = run_cell(HEAD_P, HEAD_LAM, children[ci], keep_running=True,
                    scatter_n=600); ci += 1
    print()
    print(f"{'selected so far':>16} {'running Delta':>15} {'exact':>10} {'diff':>10}")
    print(rule())
    for n_so_far, val in conv["running"]:
        print(f"{n_so_far:16,} {val:15.6f} {ha['delta']:10.6f} "
              f"{val - ha['delta']:+10.6f}")
    print()
    print(f"final running estimate : {conv['running'][-1][1]:.6f}")
    print(f"exact value            : {ha['delta']:.6f}")
    print(f"replicate-SE           : {conv['delta_se']:.6f}")
    print(f"final z                : "
          f"{zscore(conv['running'][-1][1], ha['delta'], conv['delta_se']):+.2f}")

    # -------------------------------------------------------- Figure data out
    head("FIGURE DATA")

    print("FIG1  scatter of first against second measurement, headline cell")
    print("      600 individuals drawn at random, x1 then x2, three decimals")
    print(f"      selection cut at z_p = {ha['zp']:.4f}")
    sx, sy = conv["scatter"]
    for i in range(len(sx)):
        print(f"  {sx[i]:.3f} {sy[i]:.3f}")
    print()

    print("FIG2  Delta against selection threshold p, one series per lam")
    print(f"{'lam':>6} {'r':>8} {'p':>6} {'Delta sim':>11} {'SE':>9} {'Delta exact':>12}")
    for lam in LAM_GRID:
        for p in P_GRID:
            c, a = cells[(p, lam)]
            print(f"{lam:6.2f} {a['r']:8.4f} {p:6.3f} {c['delta']:11.5f} "
                  f"{c['delta_se']:9.5f} {a['delta']:12.5f}")
    print()

    print("FIG2b analytic Delta curve, dense p, for drawing smooth lines")
    for lam in LAM_GRID:
        if lam == 0.0:
            continue
        vals = []
        for j in range(60):
            pp = 0.01 * ((0.60 / 0.01) ** (j / 59.0))
            vals.append(f"{pp:.4f}:{analytic(pp, lam)['delta']:.5f}")
        print(f"  lam={lam:.2f} " + " ".join(vals))
    print()

    print("FIG3  validation, measured shrinkage against exact r, all cells")
    print(f"{'p':>6} {'lam':>6} {'r exact':>10} {'shrink sim':>12} {'SE':>10} {'z':>8}")
    for p in P_GRID:
        for lam in LAM_GRID:
            c, a = cells[(p, lam)]
            z = zscore(c["shrink"], a["r"], c["shrink_se"])
            zs = "nan" if not math.isfinite(z) else f"{z:8.3f}"
            print(f"{p:6.2f} {lam:6.2f} {a['r']:10.6f} {c['shrink']:12.6f} "
                  f"{c['shrink_se']:10.6f} {zs:>8}")
    print()

    print("FIG4  convergence trace, headline cell (n selected, running Delta)")
    for n_so_far, val in conv["running"]:
        print(f"  {n_so_far} {val:.6f}")
    print(f"  exact {ha['delta']:.6f}   rep_se {conv['delta_se']:.6f}")
    print()

    print("FIG5  single-arm against randomised estimate, by true tau")
    for label, p, lam in TRIAL_SCENARIOS:
        a = analytic(p, lam)
        print(f"  scenario={label} p={p} lam={lam} r={a['r']:.4f} bias={a['delta']:.5f}")
        for tau in TAU_GRID:
            tr = trials[(p, lam, tau)]
            print(f"    tau={tau:.2f} single={tr['single']:.5f} "
                  f"single_se={tr['single_se']:.5f} rct={tr['rct']:.5f} "
                  f"rct_se={tr['rct_se']:.5f}")
    print()

    # ---------------------------------------------------------------- Summary
    head("SUMMARY OF CHECKS")
    print(f"1. Shrinkage law E[X2|sel] = r E[X1|sel], {n_cells} cells, "
          f"largest |z| = {max_abs_z_shrink:.2f}")
    print(f"2. Spurious improvement against (1-r) phi(z_p)/p, {n_cells} cells, "
          f"largest |z| = {max_abs_z_delta:.2f}")
    print(f"3. Zero-noise control at r = 1: "
          f"{'PASSED, Delta identically 0 at every p' if zero_ok else 'FAILED'}")
    print(f"4. Randomised design recovers the true tau, {n_trials} cells, "
          f"largest |z| = {max_abs_z_rct:.2f}")
    print(f"5. Pooled goodness of fit on Delta: sum z^2 = {chi:.2f} on {dfree} df, "
          f"p = {chi2_sf(chi, dfree):.3f}")
    worst_overall = max(max_abs_z_shrink, max_abs_z_delta, max_abs_z_rct)
    print()
    print(f"Largest |z| anywhere in the study: {worst_overall:.2f}")
    print("Checks 1 and 2 are not independent of one another: the shrinkage factor and")
    print("Delta are the same measurement written two ways, so a cell that wanders in")
    print("one wanders in the other. Counting them as one family of 42 and the trial")
    print("cells as a second family of " + str(n_trials) + ",")
    tc2 = t_crit_two_sided(0.05 / (n_cells + n_trials), N_REPS - 1)
    print(f"the Bonferroni threshold at family-wise 5% is |t| = {tc2:.2f} with "
          f"{N_REPS - 1} df.")
    if worst_overall <= tc2 and not flagged:
        print("No disagreement between simulation and closed form survives.")
    elif not flagged:
        print("One cell sits above the threshold but no cell was flagged for re-run.")
    else:
        print("See the re-run table in Part A.2 for what happened to the flagged cells.")
    print()
    print(f"headline number: selecting the bottom {HEAD_P:.0%} of a population on a")
    print(f"measurement with reliability r = {ha['r']:.2f} and applying no treatment")
    print(f"produces an apparent improvement of {hc['delta']:.4f} "
          f"+/- {hc['delta_se']:.4f} population SDs")
    print(f"(exact value {ha['delta']:.4f}).")
    print()
    print(f"wall clock: {time.time() - t_start:.1f} s")
    print(rule("="))


if __name__ == "__main__":
    main()
