#!/usr/bin/env python3
"""Hands-on P10 — an A/B testing platform, assembled from eight lego blocks."""
import hashlib, math
import numpy as np
from _harness import block, run_all

@block(1, "Assignment: deterministic, not random", "the same user must get the same arm forever")
def b1(s, show):
    def assign(uid, exp, arms=2, salt="v1"):
        h = hashlib.sha256(f"{exp}:{salt}:{uid}".encode()).digest()
        return int.from_bytes(h[:8], "big") % arms
    if show:
        N = 200_000
        a = np.array([assign(f"u{i}", "checkout") for i in range(N)])
        print(f"  {N} users hashed into 2 arms: "
              f"{[int((a==k).sum()) for k in (0,1)]}  "
              f"(imbalance {abs((a==0).mean()-.5)*100:.3f}%)")
        again = [assign(f"u{i}", "checkout") for i in range(1000)]
        print(f"  re-assigning the first 1000: identical = "
              f"{list(a[:1000]) == again}")
        b = np.array([assign(f"u{i}", "banner") for i in range(N)])
        print(f"  correlation with a SECOND experiment's assignment: "
              f"{np.corrcoef(a, b)[0,1]:+.4f}")
        print("  Deterministic hashing gives three properties at once: a returning")
        print("  user sees a consistent experience, no assignment table has to be")
        print("  stored, and two experiments are independent because the experiment")
        print("  name is inside the hash. Seeding an RNG per request gives you none")
        print("  of these -- and the bug is invisible until someone reloads a page.")
    return {"assign": assign}

@block(2, "Sample size before the test", "the number that decides whether the test is worth running")
def b2(s, show):
    def n_per_arm(p, mde_rel, alpha=0.05, power=0.8):
        z_a, z_b = 1.959964, 0.841621
        d = p * mde_rel
        return math.ceil(2 * (z_a + z_b) ** 2 * p * (1 - p) / d ** 2)
    if show:
        print(f"  baseline conversion 5%, alpha=0.05, power=0.80:")
        print(f"  {'relative MDE':>14}{'n per arm':>12}{'days @ 20k/day/arm':>22}")
        for mde in (0.20, 0.10, 0.05, 0.02, 0.01):
            n = n_per_arm(0.05, mde)
            print(f"  {mde:>13.0%}{n:>12,}{n/20_000:>22.1f}")
        print("  n scales as 1/MDE^2: detecting a 1% lift instead of a 2% one costs")
        print("  4x the traffic, not 2x. This table is the single most useful artefact")
        print("  an experimentation platform produces, because most proposed tests are")
        print("  revealed as impossible BEFORE anyone builds the feature.")
    return {"n_per_arm": n_per_arm}

@block(3, "The t-test, and what it promises", "5% false positives, by construction")
def b3(s, show):
    rng = np.random.default_rng(10)
    def welch(a, b):
        ma, mb = a.mean(), b.mean(); va, vb = a.var(ddof=1), b.var(ddof=1)
        na, nb = len(a), len(b)
        se = math.sqrt(va/na + vb/nb)
        if se == 0: return 0.0, 1.0
        t = (mb - ma) / se
        df = (va/na + vb/nb)**2 / ((va/na)**2/(na-1) + (vb/nb)**2/(nb-1))
        # normal approximation to the t CDF is fine at these df
        p = 2 * (1 - 0.5 * (1 + math.erf(abs(t) / math.sqrt(2))))
        return t, p
    if show:
        fp = 0; T = 2000; n = 4000
        for _ in range(T):
            a = rng.binomial(1, 0.05, n).astype(float)
            b = rng.binomial(1, 0.05, n).astype(float)   # A/A: NO real effect
            if welch(a, b)[1] < 0.05: fp += 1
        lo, hi = fp/T - 1.96*math.sqrt(.05*.95/T), fp/T + 1.96*math.sqrt(.05*.95/T)
        print(f"  {T} A/A tests, no effect present, n={n} per arm")
        print(f"  significant at p<0.05: {fp} ({fp/T:.3%})")
        print(f"  expected 5.000%, 95% interval [{lo:.3%}, {hi:.3%}] -> "
              f"{'calibrated' if lo <= 0.05 <= hi else 'MISCALIBRATED'}")
        print("  An A/A test is the platform's own unit test. Run a few thousand")
        print("  before you trust a single A/B result: it validates the assignment,")
        print("  the metric pipeline, and the statistics in one shot.")
    return {"welch": welch}

@block(4, "Peeking", "the most expensive statistical error in industry")
def b4(s, show):
    rng = np.random.default_rng(11)
    def trial(peeks, n=8000, p=0.05, effect=0.0):
        a = rng.binomial(1, p, n).astype(float)
        b = rng.binomial(1, p*(1+effect), n).astype(float)
        checks = np.linspace(n//peeks, n, peeks).astype(int)
        for c in checks:
            if s["welch"](a[:c], b[:c])[1] < 0.05: return True
        return False
    if show:
        T = 2000
        print(f"  A/A tests again -- no effect -- but the analyst checks the dashboard")
        print(f"  {'times checked':>15}{'false positive rate':>22}{'inflation':>12}")
        base = None
        for peeks in (1, 2, 5, 10, 50):
            fp = sum(trial(peeks) for _ in range(T)) / T
            base = base or fp
            print(f"  {peeks:>15}{fp:>21.1%}{fp/base:>12.1f}x")
        print("  The test is honest; the STOPPING RULE is not. Each look is another")
        print("  chance for noise to cross the line, and 'we stopped when it hit")
        print("  significance' converts a 5% error rate into 20%+. Fixes: fix n in")
        print("  advance and do not look, use alpha-spending, or use a sequential test")
        print("  that is valid at every moment (mSPRT, always-valid confidence")
        print("  sequences). Anything but staring at a p-value and waiting.")
    return {"trial": trial}

@block(5, "Sample ratio mismatch", "the cheapest bug detector you will ever write")
def b5(s, show):
    def srm(counts, expected=None):
        n = sum(counts); k = len(counts)
        exp = expected or [n/k]*k
        chi = sum((c-e)**2/e for c, e in zip(counts, exp))
        p = math.exp(-chi/2) if k == 2 else float("nan")   # chi2 df=1 survival
        return chi, p
    if show:
        print(f"  {'observed split':<26}{'chi2':>9}{'p':>10}{'verdict':>12}")
        for a, b, lbl in ((50_000, 50_000, "50000 / 50000"),
                          (50_000, 49_800, "50000 / 49800"),
                          (50_000, 49_400, "50000 / 49400"),
                          (50_000, 48_000, "50000 / 48000")):
            chi, p = srm([a, b])
            print(f"  {lbl:<26}{chi:>9.2f}{p:>10.2e}"
                  f"{('OK' if p > 0.001 else 'SRM -- STOP'):>12}")
        print("  A 1.2% imbalance is a 0.6% deviation per arm and looks like nothing.")
        print("  It is p<0.001 at this traffic, and it means users were lost")
        print("  NON-RANDOMLY -- a redirect that dropped slow clients, a crash in one")
        print("  arm, a bot filter that fired asymmetrically. Whatever the metric")
        print("  says afterwards is unusable, because the arms are no longer")
        print("  comparable populations. Check SRM first, always, before the metric.")
    return {"srm": srm}

@block(6, "Variance reduction with CUPED", "the same decision, on a fraction of the traffic")
def b6(s, show):
    rng = np.random.default_rng(12)
    def cuped(y, x):
        theta = np.cov(y, x)[0, 1] / np.var(x, ddof=1)
        return y - theta * (x - x.mean()), theta
    if show:
        n = 20_000
        pre = rng.gamma(2, 3, 2*n)                       # pre-period spend
        noise = rng.normal(0, 3, 2*n)
        post = 0.8 * pre + noise                         # correlated post-period
        post[n:] *= 1.02                                 # +2% true effect in arm B
        a, b = post[:n], post[n:]
        pa, pb = pre[:n], pre[n:]
        t0, p0 = s["welch"](a, b)
        ac, th = cuped(a, pa); bc, _ = cuped(b, pb)
        t1, p1 = s["welch"](ac, bc)
        r = np.corrcoef(post, pre)[0, 1]
        print(f"  correlation(pre-period, post-period) = {r:.3f}, theta = {th:.3f}")
        print(f"  {'estimator':<22}{'std error':>12}{'t':>9}{'p':>11}")
        for lbl, x, y, t, p in (("raw difference", a, b, t0, p0),
                                ("CUPED-adjusted", ac, bc, t1, p1)):
            se = math.sqrt(x.var(ddof=1)/len(x) + y.var(ddof=1)/len(y))
            print(f"  {lbl:<22}{se:>12.4f}{t:>9.2f}{p:>11.2e}")
        red = 1 - (bc.var()/b.var())
        print(f"  variance reduced {red:.1%}, which is 1 - r^2 = {1-r*r:.1%} off by")
        print(f"  {abs(red-(1-(1-r*r)))*100:.1f}pp -- the theory predicts the measurement.")
        print(f"  Equivalent traffic saving: the same power at {1-red:.0%} of n.")
        print("  CUPED is free: pre-period data already exists, and the adjustment")
        print("  cannot bias the estimate because x is measured BEFORE assignment.")
    return {"cuped": cuped}

@block(7, "Multiple metrics, multiple arms", "twenty metrics guarantee a winner")
def b7(s, show):
    rng = np.random.default_rng(13)
    if show:
        T, M, n = 1000, 20, 5000
        any_sig = bh_sig = bonf_sig = per_metric = 0
        for _ in range(T):
            ps = []
            for _ in range(M):
                a = rng.binomial(1, .05, n).astype(float)
                b = rng.binomial(1, .05, n).astype(float)
                ps.append(s["welch"](a, b)[1])
            ps = np.sort(np.array(ps))
            per_metric += int((ps < 0.05).sum())
            any_sig += ps[0] < 0.05
            bonf_sig += ps[0] < 0.05 / M
            bh = ps <= 0.05 * np.arange(1, M+1) / M       # Benjamini-Hochberg
            bh_sig += bh.any()
        print(f"  {T} A/A experiments, {M} metrics each, no effect anywhere:")
        print(f"  {'rule':<34}{'experiments with a winner':>28}")
        for lbl, v in (("any metric p<0.05 (no correction)", any_sig),
                       ("Bonferroni (p < 0.05/20)", bonf_sig),
                       ("Benjamini-Hochberg FDR 5%", bh_sig)):
            print(f"  {lbl:<34}{v/T:>27.1%}")
        r = per_metric / (T * M)
        print(f"  Textbook: 1-(1-0.05)^20 = {1-0.95**20:.1%}. Measured {any_sig/T:.1%}.")
        print(f"  The gap is not sampling noise -- it is that the union bound needs the")
        print(f"  ACTUAL per-metric rate, which was {r:.2%} here, not the nominal 5%")
        print(f"  (the normal approximation to Welch's t is mildly anti-conservative on")
        print(f"  binary data). 1-(1-{r:.4f})^20 = {1-(1-r)**20:.1%}, which matches.")
        print(f"  A 0.4pp error per metric compounds into a 3.5pp error across twenty.")
        print("  Declare ONE primary metric before the test. Everything else is a")
        print("  guardrail (checked for harm, one-sided) or exploratory (reported,")
        print("  never used to declare a win). This is a process rule, not a")
        print("  statistical one -- which is why the platform should enforce it.")
    return {}

@block(8, "Power, honestly", "an underpowered test is worse than no test")
def b8(s, show):
    rng = np.random.default_rng(14)
    if show:
        print("  A REAL +5% relative effect exists. How often do we find it, and what")
        print("  does the estimate look like when we do?")
        print(f"  {'n per arm':>11}{'power':>9}{'mean lift | significant':>26}"
              f"{'exaggeration':>14}")
        for n in (2_000, 10_000, 30_000, 120_000):
            hits, ests = 0, []
            for _ in range(600):
                a = rng.binomial(1, .05, n).astype(float)
                b = rng.binomial(1, .0525, n).astype(float)
                t, p = s["welch"](a, b)
                if p < 0.05 and b.mean() > a.mean():
                    hits += 1; ests.append((b.mean()-a.mean())/a.mean())
            m = float(np.mean(ests)) if ests else float("nan")
            print(f"  {n:>11,}{hits/600:>9.1%}{m:>25.1%}{m/0.05:>13.1f}x")
        print("  This is the type-M (magnitude) error. An underpowered test does not")
        print("  just miss effects -- when it DOES find one, the estimate is inflated,")
        print("  because only the luckiest samples clear the threshold. Shipping on a")
        print("  20%-powered test means the launch report overstates the win ~2x, and")
        print("  the follow-up 'why did the metric not move in production' is")
        print("  guaranteed. Compute power before, not after.")
    return {}

def assembly(s):
    print("\nEight blocks = an experimentation platform. One experiment, end to end.\n")
    rng = np.random.default_rng(20)
    N, TRUE = 60_000, 0.03
    uids = [f"user{i}" for i in range(N)]
    arm = np.array([s["assign"](u, "checkout-redesign") for u in uids])
    pre = rng.gamma(2, 3, N)
    conv = rng.binomial(1, np.where(arm == 1, .05*(1+TRUE), .05)).astype(float)
    rev = conv * (0.7*pre + rng.normal(0, 2, N)) 

    print("  STEP 1  design")
    need = s["n_per_arm"](0.05, 0.03)
    print(f"    to detect a {TRUE:.0%} relative lift at 80% power: "
          f"{need:,} per arm; we have {int((arm==0).sum()):,}")
    print(f"    -> the test is {'ADEQUATELY POWERED' if (arm==0).sum() >= need else 'UNDERPOWERED, and we run it anyway to see what that looks like'}")

    print("  STEP 2  health checks")
    chi, p = s["srm"]([int((arm==0).sum()), int((arm==1).sum())])
    print(f"    SRM: {int((arm==0).sum())} / {int((arm==1).sum())}  "
          f"chi2={chi:.2f}  p={p:.3f}  -> {'PASS' if p > 0.001 else 'FAIL'}")

    print("  STEP 3  primary metric, fixed horizon, no peeking")
    a, b = conv[arm == 0], conv[arm == 1]
    t, pv = s["welch"](a, b)
    lift = (b.mean()-a.mean())/a.mean()
    print(f"    conversion  A={a.mean():.4f}  B={b.mean():.4f}  "
          f"lift={lift:+.2%}  p={pv:.4f}")
    print(f"    true lift was {TRUE:+.0%}; the estimate is "
          f"{'inside' if abs(lift-TRUE) < 2*math.sqrt(a.var()/len(a)+b.var()/len(b))/a.mean() else 'outside'}"
          " a 2-SE window of it")

    print("  STEP 4  the same metric, CUPED-adjusted")
    ra, rb = rev[arm == 0], rev[arm == 1]
    ca, _ = s["cuped"](ra, pre[arm == 0]); cb, _ = s["cuped"](rb, pre[arm == 1])
    t2, p2 = s["welch"](ra, rb); t3, p3 = s["welch"](ca, cb)
    print(f"    revenue/user  raw   p={p2:.4f}   se="
          f"{math.sqrt(ra.var(ddof=1)/len(ra)+rb.var(ddof=1)/len(rb)):.4f}")
    print(f"    revenue/user  CUPED p={p3:.4f}   se="
          f"{math.sqrt(ca.var(ddof=1)/len(ca)+cb.var(ddof=1)/len(cb)):.4f}")
    rr = np.corrcoef(rev, pre)[0, 1]
    print(f"    barely moved, and block 6 says exactly why: the gain is 1-r^2 and")
    print(f"    here r={rr:.3f}, so the ceiling is {1-(1-rr*rr):.1%}. Revenue/user is")
    print(f"    zero-inflated -- {100*(rev==0).mean():.0f}% of users never convert -- so a")
    print(f"    pre-period covariate cannot explain much of it. CUPED is not a free")
    print(f"    win; it is a free win ON METRICS THAT AUTOCORRELATE. Check r first.")

    print("  STEP 5  what peeking would have done to this test")
    hit = [c for c in range(2000, N//2, 2000)
           if s["welch"](a[:c], b[:c])[1] < 0.05]
    nchk = len(range(2000, N//2, 2000))
    print(f"    p<0.05 at {len(hit)} of {nchk} checkpoints; "
          f"first at n={hit[0] if hit else '--'}")
    print(f"    final verdict at the pre-registered n: p={pv:.4f}")
    print(f"    This run got away with it. Block 4 is the reason that is luck and not")
    print(f"    method: at {nchk} looks the false-positive rate is ~20%, so one launch")
    print(f"    in five would have shipped a null result as a win. A single experiment")
    print(f"    can never tell you whether your process is sound -- only the")
    print(f"    distribution over many can, which is what blocks 3, 4, 7 and 8 do.")

    print("\n  Every block appears in that sequence, in the order a real launch uses")
    print("  it: power first (or do not run), health checks second (or do not read),")
    print("  one primary metric third, variance reduction fourth, and the peeking")
    print("  analysis as a reminder of what the other path looked like.")
    print("\n  The platform's value is NOT the t-test -- that is twelve lines in")
    print("  block 3. It is that the sequence above is enforced by software instead")
    print("  of remembered by people under launch pressure.")
    print("\n  Built: hash assignment -> power -> t-test + A/A calibration -> peeking")
    print("  -> SRM -> CUPED -> multiple comparisons -> type-M error.")
    print("  Missing, on the project page: metric definition and a metrics repo (m3),")
    print("  the delta method for ratio metrics with a user-level denominator (m7),")
    print("  switchback and cluster randomisation for interference (m9), sequential")
    print("  tests with always-valid intervals (m10), and E8 -- the heterogeneous")
    print("  treatment effect analysis that finds the segment the average hides.")

if __name__ == "__main__":
    run_all(assembly, "HANDS-ON P10 — A/B testing platform, block by block")