# %% [markdown]
# # Power for a four-arm savings-reminder trial
#
# **Worked example on synthetic data, by Ahmed Beddi.** No real programme, institution or dataset is
# involved. Parameters (baseline rate, effect size, attrition) are planning assumptions chosen for
# illustration and are stated where they are set.
#
# The notebook does two things:
#
# 1. computes the sample size analytically, with the same formulas as the `PowerCalculator` demo on the
#    website (same defaults, same rounding);
# 2. checks that number by simulating the trial many times and running the **pre-specified analysis
#    model** (linear probability model with stratum fixed effects, HC2 robust standard errors, Holm
#    correction across the three treatment–control contrasts) on each simulated dataset.
#
# Design: individual randomisation of active mobile-wallet holders to four arms (control, plain
# reminder, reminder + goal framing, reminder + loss framing), stratified on sex and baseline activity.
# Primary outcome: *regular saver*, an indicator equal to 1 if the person made a savings deposit in at
# least 3 of the 6 months after randomisation, measured from administrative transaction records.

# %%
import numpy as np
from math import ceil, sqrt
from scipy.stats import norm
import matplotlib.pyplot as plt

# %% [markdown]
# ## 1. Analytic sample size
#
# For each treatment arm $k$ compared with control, with control-to-treatment allocation ratio $r$
# ($n_c = r\,n_t$), significance level $\alpha^\ast$ after multiplicity adjustment and power $1-\beta$:
#
# $$n_t^{\text{raw}} = \frac{(z_{1-\alpha^\ast/2} + z_{1-\beta})^2\,V}{\delta^2},\qquad
# V = \begin{cases} p_1(1-p_1) + p_0(1-p_0)/r & \text{binary outcome} \\ \sigma^2\,(1 + 1/r) & \text{continuous outcome}\end{cases}$$
#
# with $p_1 = p_0 + \delta$. We then inflate by the design effect $\text{DE} = 1 + (m-1)\rho$ (cluster size
# $m$, intra-cluster correlation $\rho$) and by expected attrition $a$:
#
# $$n_t = \left\lceil \frac{n_t^{\text{raw}}\cdot \text{DE}}{1-a} \right\rceil,\qquad n_c = \lceil r\,n_t\rceil,\qquad N = (A-1)\,n_t + n_c .$$
#
# With $A$ arms there are $K = A-1$ confirmatory contrasts against control. For planning we use
# $\alpha^\ast = \alpha/K$ for both Bonferroni and Holm: Holm's first step uses exactly $\alpha/K$, and every
# later step is more lenient, so a design powered at $\alpha/K$ has *at least* the stated power for each
# contrast under Holm.

# %%
def plan(outcome='binary', p0=0.30, delta=0.05, sd=1.0, alpha=0.05, power=0.80, arms=4,
         mc='holm', m=1, icc=0.0, attrition=0.05, r=1.0):
    K = max(arms - 1, 1)
    a_star = alpha / K if (mc != 'none' and K > 1) else alpha
    za, zb = norm.ppf(1 - a_star / 2), norm.ppf(power)
    if outcome == 'binary':
        p1 = p0 + delta
        V = p1 * (1 - p1) + p0 * (1 - p0) / r
    else:
        V = sd ** 2 * (1 + 1 / r)
    n_raw = (za + zb) ** 2 * V / delta ** 2
    de = 1 + (m - 1) * icc
    n_t = ceil(n_raw * de / (1 - attrition))
    n_c = ceil(r * n_t)
    return dict(alpha_star=a_star, z_alpha=za, z_beta=zb, V=V, n_raw=n_raw, DE=de,
                n_t=n_t, n_c=n_c, N=(arms - 1) * n_t + n_c,
                clusters_t=ceil(n_t / m), clusters_c=ceil(n_c / m))

def power_at(n_t, outcome='binary', p0=0.30, delta=0.05, sd=1.0, alpha=0.05, arms=4, mc='holm',
             m=1, icc=0.0, attrition=0.05, r=1.0):
    """Analytic power for one contrast, given enrolled treated-arm size n_t (same approximation as the demo)."""
    K = max(arms - 1, 1)
    a_star = alpha / K if (mc != 'none' and K > 1) else alpha
    za = norm.ppf(1 - a_star / 2)
    de = 1 + (m - 1) * icc
    nt_eff = n_t * (1 - attrition) / de
    nc_eff = r * n_t * (1 - attrition) / de
    if outcome == 'binary':
        p1 = p0 + delta
        se = sqrt(p1 * (1 - p1) / nt_eff + p0 * (1 - p0) / nc_eff)
    else:
        se = sd * sqrt(1 / nt_eff + 1 / nc_eff)
    return norm.cdf(abs(delta) / se - za)

P = plan()
for k, v in P.items():
    print(f'{k:>11}: {v:.4f}' if isinstance(v, float) else f'{k:>11}: {v}')

# %% [markdown]
# **Headline (defaults):** baseline regular-saver rate 30 %, minimum detectable effect 5 percentage points,
# $\alpha = 0.05$ two-sided split across three contrasts, power 80 %, 5 % attrition, equal allocation.
# The printout above gives the treated-arm size $n_t$ and the total $N$ for four arms; these are the
# numbers the website demo shows with its default settings.
#
# ### Sensitivity
# The table shows how the requirement moves with the two assumptions that matter most: the baseline
# rate and the effect we insist on detecting.

# %%
import pandas as pd
rows = []
for p0 in [0.20, 0.30, 0.40]:
    for d in [0.03, 0.04, 0.05, 0.07]:
        q = plan(p0=p0, delta=d)
        rows.append(dict(baseline=p0, MDE_pp=round(100 * d), n_per_arm=q['n_t'], total_N=q['N']))
sens = pd.DataFrame(rows).pivot(index='MDE_pp', columns='baseline', values='total_N')
print('Total N (4 arms, Holm planning alpha, 5% attrition)')
print(sens)

# %% [markdown]
# Two further levers, for reference: without a multiplicity adjustment the requirement would fall
# (but the trial would no longer control the family-wise error rate), and allocating more people to
# control ($r = \sqrt{K} \approx 1.73$, the classic choice when every arm is compared with one control)
# lowers total $N$ for the same per-contrast power.

# %%
for label, kw in [('no adjustment', dict(mc='none')), ('Holm, r = 1', {}), ('Holm, r = sqrt(3)', dict(r=sqrt(3)))]:
    q = plan(**kw)
    print(f'{label:<20} n_t = {q["n_t"]:>5}  n_c = {q["n_c"]:>5}  N = {q["N"]:>5}')

# %% [markdown]
# ### Power curve (same curve as the demo)

# %%
ns = np.linspace(200, 2 * P['n_t'], 120)
pw = [power_at(n) for n in ns]
fig, ax = plt.subplots(figsize=(6.5, 3.6))
ax.plot(ns, pw, color='#a3471c')
ax.axhline(0.8, ls='--', color='grey', lw=0.8)
ax.axvline(P['n_t'], ls=':', color='#1f5f63')
ax.set_xlabel('Enrolled per treatment arm'); ax.set_ylabel('Power, one contrast')
ax.set_ylim(0, 1); ax.set_title('Power vs sample size (p0 = 0.30, MDE = 5 pp, Holm)')
plt.tight_layout(); plt.show()
print('power at n_t:', round(power_at(P['n_t']), 4))

# %% [markdown]
# ## 2. Power by simulation of the pre-specified analysis model
#
# The analytic formula is an approximation. The test that matters is whether the *analysis we will
# actually run* rejects as often as planned. For each replication we:
#
# 1. create $N$ participants in four strata (sex × baseline activity), randomised in equal shares to the
#    four arms within each stratum;
# 2. drop 5 % at random (closed accounts, ported numbers: outcome missing);
# 3. draw the outcome; 4. estimate
#
# $$Y_i = \alpha + \sum_{k=1}^{3} \tau_k T_{ki} + \sum_{s=2}^{4} \gamma_s S_{si} + \varepsilon_i$$
#
# by OLS with HC2 standard errors; 5. compute two-sided p-values for $\tau_1,\tau_2,\tau_3$ and apply
# Holm (and, for comparison, Bonferroni) at family-wise $\alpha = 0.05$.
#
# The OLS/HC2 step is written in numpy for speed; the next cell checks it against statsmodels.

# %%
def ols_hc2(y, X):
    XtX_inv = np.linalg.inv(X.T @ X)
    b = XtX_inv @ X.T @ y
    e = y - X @ b
    h = np.einsum('ij,jk,ik->i', X, XtX_inv, X)
    meat = (X * (e ** 2 / (1 - h))[:, None]).T @ X
    V = XtX_inv @ meat @ XtX_inv
    return b, np.sqrt(np.diag(V))

def holm_reject(p, alpha=0.05):
    order = np.argsort(p); K = len(p); rej = np.zeros(K, bool)
    for step, idx in enumerate(order):
        if p[idx] <= alpha / (K - step):
            rej[idx] = True
        else:
            break
    return rej

def simulate_once(rng, n_t, effects, p_strata=(0.30, 0.30, 0.30, 0.30), attrition=0.05):
    arms = len(effects) + 1
    n = n_t * arms
    strata = rng.integers(0, 4, n)
    arm = np.empty(n, int)
    for s in range(4):                       # complete randomisation within stratum
        idx = np.flatnonzero(strata == s)
        lab = np.resize(np.arange(arms), idx.size)
        arm[idx] = rng.permutation(lab)
    keep = rng.random(n) >= attrition
    strata, arm = strata[keep], arm[keep]
    p = np.asarray(p_strata)[strata] + np.r_[0.0, effects][arm]
    y = (rng.random(arm.size) < p).astype(float)
    X = np.column_stack([np.ones(arm.size)] + [(arm == k).astype(float) for k in range(1, arms)]
                        + [(strata == s).astype(float) for s in range(1, 4)])
    b, se = ols_hc2(y, X)
    t = b[1:arms] / se[1:arms]
    pval = 2 * norm.sf(np.abs(t))
    return pval, (y, X)

# cross-check one replication against statsmodels
import statsmodels.api as sm
rng = np.random.default_rng(1)
pval, (y, X) = simulate_once(rng, P['n_t'], [0.05, 0.05, 0.05])
fit = sm.OLS(y, X).fit(cov_type='HC2')
b, se = ols_hc2(y, X)
print('max |coef diff|:', np.abs(fit.params - b).max(), ' max |SE diff|:', np.abs(fit.bse - se).max())

# %%
def power_sim(effects, reps=2000, seed=20261103, **kw):
    rng = np.random.default_rng(seed)
    holm = np.zeros((reps, 3), bool); bonf = np.zeros((reps, 3), bool)
    for i in range(reps):
        pval, _ = simulate_once(rng, P['n_t'], effects, **kw)
        holm[i] = holm_reject(pval); bonf[i] = pval <= 0.05 / 3
    return holm, bonf

scenarios = {
    'A. all three reminders +5 pp': [0.05, 0.05, 0.05],
    'B. only goal framing +5 pp':   [0.00, 0.05, 0.00],
    'C. global null (no effects)':  [0.00, 0.00, 0.00],
}
res = {}
for name, eff in scenarios.items():
    holm, bonf = power_sim(eff)
    res[name] = (holm, bonf)
    print(name)
    print('   rejection rate per contrast, Holm      :', holm.mean(0).round(3))
    print('   rejection rate per contrast, Bonferroni:', bonf.mean(0).round(3))
    print('   any false rejection (FWER), Holm       :',
          round((holm & (np.array(eff) == 0)).any(1).mean(), 3))
print('analytic power per contrast:', round(power_at(P['n_t']), 3), ' Monte Carlo SE ~', round(sqrt(0.8 * 0.2 / 2000), 3))

# %% [markdown]
# **Reading the simulation.**
#
# * Scenario B is the planning scenario (one framing works, the others do not). Simulated power for
#   that contrast should sit at about the analytic 80 %, within Monte Carlo error (about ±2 points at
#   two standard errors). Holm and Bonferroni give almost the same power here: with only one real
#   effect, Holm's first (and binding) threshold is the Bonferroni threshold.
# * Scenario A shows that Holm's power per contrast rises when several arms work: after the first
#   rejection the remaining thresholds relax to $\alpha/2$ and then $\alpha$.
# * Scenario C checks the error rate we promise: the probability of declaring *any* reminder effective
#   when none is should be at or below 5 %.
#
# ## 3. Covariate gain: why the analytic number is conservative
#
# In the scenarios above the strata are not prognostic, so stratum fixed effects add nothing. In a real
# wallet base, people already active at baseline are far more likely to be regular savers. If the
# regular-saver rate is, say, 10 % among the baseline-dormant and 50 % among the baseline-active (still
# 30 % on average), the same design and the same model have more power. We do not plan on this gain;
# it is a margin of safety against an optimistic effect size.

# %%
holm, bonf = power_sim([0.00, 0.05, 0.00], p_strata=(0.10, 0.10, 0.50, 0.50))
print('Scenario B with prognostic strata, Holm power for goal-framing contrast:', holm.mean(0)[1].round(3))

# %% [markdown]
# ## 4. Minimum detectable effect for a fixed budget
#
# If the partner can only include, say, 1,200 people per arm, what is the smallest effect the trial can
# reliably detect? Solve $\text{power}(\delta) = 0.8$ for $\delta$ by bisection (the demo does the same).

# %%
def mde_at(n_t, target=0.80, **kw):
    lo, hi = 1e-6, 1 - kw.get('p0', 0.30) - 1e-6
    for _ in range(80):
        mid = (lo + hi) / 2
        if power_at(n_t, delta=mid, **kw) < target:
            lo = mid
        else:
            hi = mid
    return hi

for n in [800, 1200, P['n_t'], 2500]:
    print(f'n_t = {n:>5}:  MDE = {100 * mde_at(n):.2f} percentage points')

# %% [markdown]
# ## 5. What the trial is *not* powered for: differences between framings
#
# Comparing two active framings (goal vs loss) is a secondary question. If both raise the regular-saver
# rate and differ by 2 percentage points (35 % vs 37 %), the per-arm sample needed for 80 % power at the
# same planning α is far larger than the trial. The pre-analysis plan therefore reports framing
# contrasts as secondary and does not let the programme decision rest on them.

# %%
q = plan(p0=0.35, delta=0.02)
print('per arm for a 2 pp framing difference:', q['n_t'], ' vs planned per arm:', P['n_t'],
      f' (ratio {q["n_t"] / P["n_t"]:.1f}x)')