# Sample size for a two-arm parallel trial. Primary outcome: the directed-exploration parameter
# (information bonus) fitted to the Horizon task (Wilson et al. 2014), one week after the course,
# analysed by ANCOVA adjusting for the baseline value. Effect expressed in baseline SD units.
import math
from statistics import NormalDist
N = NormalDist()

def n_per_arm(d, r=0.5, alpha=0.05, power=0.90):
    """Normal approximation; ANCOVA on a baseline correlated r with the outcome shrinks the
    variance by (1 - r^2)."""
    za, zb = N.inv_cdf(1 - alpha / 2), N.inv_cdf(power)
    return math.ceil(2 * (za + zb) ** 2 * (1 - r ** 2) / d ** 2) + 1

def power(n, d, r=0.5, alpha=0.05):
    za = N.inv_cdf(1 - alpha / 2)
    return N.cdf(d / math.sqrt(2 * (1 - r ** 2) / n) - za)

for d in (0.3, 0.4, 0.5):
    print(f"standardised effect {d}: {n_per_arm(d, power=0.80)} per arm for 80 % power, {n_per_arm(d)} for 90 %"
          f" (baseline-outcome correlation 0.5)")
print(f"if the correlation is only 0.3, effect 0.4 needs {n_per_arm(0.4, r=0.3)} per arm for 90 %")
# Both courses are taught in groups of 12, so outcomes within a group are correlated (ICC assumed 0.02).
de = 1 + (12 - 1) * 0.02
n = math.ceil(n_per_arm(0.4) * de)
print(f"design effect for groups of 12, ICC 0.02: {de:.2f} -> {n} completers per arm "
      f"(power without clustering would be {power(n, 0.4):.0%}); recruit {math.ceil(n / 0.85)} per arm allowing 15 % loss")
# standardised effect 0.3: 132 per arm for 80 % power, 177 for 90 % (baseline-outcome correlation 0.5)
# standardised effect 0.4: 75 per arm for 80 % power, 100 for 90 % (baseline-outcome correlation 0.5)
# standardised effect 0.5: 49 per arm for 80 % power, 65 for 90 % (baseline-outcome correlation 0.5)
# if the correlation is only 0.3, effect 0.4 needs 121 per arm for 90 %
# design effect for groups of 12, ICC 0.02: 1.22 -> 122 completers per arm (power without clustering would be 95%); recruit 144 per arm allowing 15 % loss

SESOI = 0.4     # smallest effect worth acting on, in baseline SD units of the directed-exploration parameter

def verdict(d, lo, hi):
    """d = sandbox-workshop arm minus reading-lecture arm (ANCOVA, standardised) with its 95 % CI.
    Positive = more directed exploration. Fixed before any data are seen."""
    if lo >= SESOI:
        v = "more directed exploration, at least the smallest worthwhile effect"
    elif lo > 0:
        v = "more directed exploration; size not shown to reach 0.4"
    elif hi < 0:
        v = "less directed exploration than the reading-lecture course"
    elif lo >= -SESOI and hi <= SESOI:
        v = "no difference worth acting on"
    else:
        v = "unclear"
    return f"effect {d:+.2f} ({lo:+.2f} to {hi:+.2f}): {v}"

# Invented results, to show in advance how the trial will be reported.
for res in ((0.15, -0.13, 0.43), (0.62, 0.42, 0.82), (0.05, -0.20, 0.30)):
    print(verdict(*res))
print("Secondary outcomes (random-exploration parameter, Iowa Gambling Task net score,")
print("task-switching cost) are reported with intervals but cannot make the trial positive.")
# effect +0.15 (-0.13 to +0.43): unclear
# effect +0.62 (+0.42 to +0.82): more directed exploration, at least the smallest worthwhile effect
# effect +0.05 (-0.20 to +0.30): no difference worth acting on
# Secondary outcomes (random-exploration parameter, Iowa Gambling Task net score,
# task-switching cost) are reported with intervals but cannot make the trial positive.